remi-or

#!/usr/bin/env python3
"""Visualize continuous batching timing data as a timeline."""

import json
import sys

import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

remi-or

"""
k_cache shape: torch.Size([1616, 256, 8, 128])
v_cache shape: torch.Size([1616, 256, 8, 128])
k shape: torch.Size([5, 1, 8, 128])
v shape: torch.Size([5, 1, 8, 128])
cache_seqlens shape: torch.Size([5])
block_table shape: torch.Size([1, 5, 64])
flash_kwargs: {'block_table': tensor([[ 0, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
         -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
         -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,

remi-or

from typing import Tuple
import torch
from torch import Tensor
from torchvision.transforms.v2 import functional as F


def reference_resize(image: Tensor, new_size: Tuple[int, int]) -> Tensor:
    image = F.resize(image, new_size, interpolation=F.InterpolationMode.BICUBIC, antialias=True)
    return image

remi-or

import torch
from torch import nn
from torch.nn import functional as F
from torchvision.transforms.v2 import functional as F
import matplotlib.pyplot as plt


INPUT_SIZE = (3, 41, 70)

remi-or

    def aug_test(self, 
                 imgs: List[Tensor], 
                 img_metas: List[dict],
                 rescale: bool = False) -> Tensor:
        acc_boxes = np.zeros((0, 5))
        acc_score = np.zeros((0, self.roi_head.bbox_head.num_classes))
        for img, img_meta in zip(imgs, img_metas):
            for label, dets in enumerate(self.simple_test(img, img_meta, None, rescale)[0]):
                boxes, scores = dets[:, :-1], dets[:, -1]
                acc_boxes = np.vstack((acc_boxes, boxes))

remi-or

remi-or

## Imports
from typing import Tuple
import torch
from torch import Module, Tensor
from transformers.models.roberta.modeling_roberta import RobertaPreTrainedModel, RobertaConfig, RobertaModel, RobertaEncoder
from torch.nn import CrossEntropyLoss, CosineEmbeddingLoss



## Function

remi-or

import torch
from torch.nn import CrossEntropyLoss, CosineEmbeddingLoss

def distillation_loss(
    teacher_logits : Tensor,
    student_logits : Tensor,
    labels : Tensor,
    temperature : float = 1.0,
) -> Tensor:
    """

remi-or

from torch import Tensor

def get_logits(
    model : RobertaPreTrainedModel, 
    input_ids : Tensor,
    attention_mask : Tensor,
) -> Tensor:
    """
    Given a RoBERTa (model) for classification and the couple of (input_ids) and (attention_mask),
    returns the logits corresponding to the prediction.

remi-or

from transformers.models.roberta.modeling_roberta import RobertaEncoder, RobertaModel
from torch.nn import Module

def distill_roberta_weights(
    teacher : Module,
    student : Module,
) -> None:
    """
    Recursively copies the weights of the (teacher) to the (student).
    This function is meant to be first called on a RobertaFor... model, but is then called on every children of that model recursively.