| import types |
| from typing import Optional, List, Union, Callable |
|
|
| import torch |
| from torch import nn, Tensor |
| from torch.nn import functional as F |
|
|
| from torchvision.models.mobilenetv2 import MobileNetV2 |
| from torchvision.models.resnet import ResNet |
| from torchvision.models.efficientnet import EfficientNet |
| from torchvision.models.vision_transformer import VisionTransformer |
|
|
|
|
| def compute_policy_loss(loss_sequence, mask_sequence, rewards): |
| losses = sum(mask * padded_loss for mask, padded_loss in zip(mask_sequence, loss_sequence)) |
| returns = sum(padded_reward * mask for padded_reward, mask in zip(rewards, mask_sequence)) |
| loss = torch.mean(losses * returns) |
|
|
| return loss |
|
|
|
|
| class TPBlock(nn.Module): |
| def __init__(self, depths: int, in_planes: int, out_planes: int = None, rank=1, shape_dims=3, channel_first=True, dtype=torch.float32) -> None: |
| super().__init__() |
| out_planes = in_planes if out_planes is None else out_planes |
| self.layers = torch.nn.ModuleList([self._make_layer(in_planes, out_planes, rank, shape_dims, channel_first, dtype) for _ in range(depths)]) |
|
|
| def forward(self, x: Tensor) -> Tensor: |
| for layer in self.layers: |
| x = x + layer(x) |
| return x |
|
|
| def _make_layer(self, in_planes: int, out_planes: int = None, rank=1, shape_dims=3, channel_first=True, dtype=torch.float32) -> nn.Sequential: |
| |
| class Permute(nn.Module): |
| def __init__(self, *dims): |
| super().__init__() |
| self.dims = dims |
| def forward(self, x): |
| return x.permute(*self.dims) |
| |
| class RMSNorm(nn.Module): |
| __constants__ = ["eps"] |
| eps: float |
|
|
| def __init__(self, hidden_size, eps: float = 1e-6, device=None, dtype=None): |
| """ |
| LlamaRMSNorm is equivalent to T5LayerNorm. |
| """ |
| factory_kwargs = {"device": device, "dtype": dtype} |
| super().__init__() |
| self.eps = eps |
| self.weight = nn.Parameter(torch.ones(hidden_size, **factory_kwargs)) |
|
|
| def forward(self, hidden_states): |
| input_dtype = hidden_states.dtype |
| hidden_states = hidden_states.to(torch.float32) |
| variance = hidden_states.pow(2).mean(dim=1, keepdim=True) |
| hidden_states = hidden_states * torch.rsqrt(variance + self.eps) |
| weight = self.weight.view(1, -1, *[1] * (hidden_states.ndim - 2)) |
| return weight * hidden_states.to(input_dtype) |
|
|
| def extra_repr(self): |
| return f"{self.weight.shape[0]}, eps={self.eps}" |
|
|
| conv_map = { |
| 2: (nn.Conv1d, (0, 2, 1), (0, 2, 1)), |
| 3: (nn.Conv2d, (0, 3, 1, 2), (0, 2, 3, 1)), |
| 4: (nn.Conv3d, (0, 4, 1, 2, 3), (0, 2, 3, 4, 1)), |
| } |
| Conv, pre_dims, post_dims = conv_map[shape_dims] |
| kernel_size, dilation, padding = self.generate_hyperparameters(rank) |
| |
| pre_permute = nn.Identity() if channel_first else Permute(*pre_dims) |
| post_permute = nn.Identity() if channel_first else Permute(*post_dims) |
| conv1 = Conv(in_planes, out_planes, kernel_size, padding=padding, dilation=dilation, bias=False, dtype=dtype, device='cuda') |
| nn.init.zeros_(conv1.weight) |
| bn1 = RMSNorm(out_planes, dtype=dtype, device="cuda") |
| relu = nn.ReLU(inplace=True) |
| conv2 = Conv(out_planes, in_planes, kernel_size, padding=padding, dilation=dilation, bias=False, dtype=dtype, device='cuda') |
| nn.init.zeros_(conv2.weight) |
| bn2 = RMSNorm(in_planes, dtype=dtype, device="cuda") |
|
|
| return torch.nn.Sequential(pre_permute, conv1, bn1, relu, conv2, bn2, relu, post_permute) |
|
|
| @staticmethod |
| def generate_hyperparameters(rank: int): |
| """ |
| Generates kernel size and dilation rate pairs sorted by increasing padded kernel size. |
| |
| Args: |
| rank: Number of (kernel_size, dilation) pairs to generate. Must be positive. |
| |
| Returns: |
| Tuple[int, int]: A (kernel_size, dilation) tuple where: |
| - kernel_size: Always odd and >= 1 |
| - dilation: Computed to maintain consistent padded kernel size growth |
| |
| Note: |
| Padded kernel size is calculated as: |
| (kernel_size - 1) * dilation + 1 |
| Pairs are generated first in order of increasing padded kernel size, |
| then by increasing kernel size for equal padded kernel sizes. |
| """ |
| pairs = [(1, 1, 0)] |
| padded_kernel_size = 3 |
| |
| while len(pairs) < rank: |
| for kernel_size in range(3, padded_kernel_size + 1, 2): |
| if (padded_kernel_size - 1) % (kernel_size - 1) == 0: |
| dilation = (padded_kernel_size - 1) // (kernel_size - 1) |
| padding = dilation * (kernel_size - 1) // 2 |
| pairs.append((kernel_size, dilation, padding)) |
| if len(pairs) >= rank: |
| break |
| |
| |
| padded_kernel_size += 2 |
| |
| return pairs[-1] |
|
|
|
|
| class ResNetConfig: |
| @staticmethod |
| def gen_shared_head(self): |
| def func(hidden_states): |
| """ |
| Args: |
| hidden_states (Tensor): Hidden States tensor of shape [B, C, H, W]. |
| |
| Returns: |
| logits (Tensor): Logits tensor of shape [B, C]. |
| """ |
| x = self.avgpool(hidden_states) |
| x = torch.flatten(x, 1) |
| logits = self.fc(x) |
| return logits |
| return func |
| |
| @staticmethod |
| def gen_logits(self, shared_head): |
| def func(hidden_states): |
| """ |
| Args: |
| hidden_states (Tensor): Hidden States tensor of shape [B, L, hidden_units]. |
| |
| Returns: |
| logits_seqence (List[Tensor]): List of Logits tensors. |
| """ |
| logits_sequence = [shared_head(hidden_states)] |
| for layer in self.trp_blocks: |
| logits_sequence.append(shared_head(layer(hidden_states))) |
| return logits_sequence |
| return func |
|
|
| @staticmethod |
| def gen_mask(label_smoothing=0.0, top_k=1): |
| def func(logits_sequence, labels): |
| """ |
| Args: |
| logits_sequence (List[Tensor]): List of Logits tensors. |
| labels (Tensor): Target labels of shape [B] or [B, C]. |
| |
| Returns: |
| mask_sequence (List[Tensor]): List of Mask tensor. |
| returns (Tensor): Boolean mask tensor of shape [B*(L-1)]. |
| """ |
| labels = torch.argmax(labels, dim=1) if label_smoothing > 0.0 else labels |
|
|
| mask_sequence = [torch.ones_like(labels, dtype=torch.float32, device=labels.device)] |
| for logits in logits_sequence: |
| with torch.no_grad(): |
| topk_values, topk_indices = torch.topk(logits, top_k, dim=-1) |
| mask = torch.eq(topk_indices, labels[:, None]).any(dim=-1).to(torch.float32) |
| mask_sequence.append(mask_sequence[-1] * mask) |
| return mask_sequence |
| return func |
| |
| @staticmethod |
| def gen_criterion(label_smoothing=0.0): |
| def func(logits_sequence, labels): |
| """ |
| Args: |
| logits_sequence (List[Tensor]): List of Logits tensor. |
| labels (Tensor): labels labels of shape [B] or [B, C]. |
| |
| Returns: |
| loss (Tensor): Scalar tensor representing the loss. |
| mask (Tensor): Boolean mask tensor of shape [B]. |
| """ |
| labels = torch.argmax(labels, dim=1) if label_smoothing > 0.0 else labels |
|
|
| loss_sequence = [] |
| for logits in logits_sequence: |
| loss_sequence.append(F.cross_entropy(logits, labels, reduction="none", label_smoothing=label_smoothing)) |
|
|
| return loss_sequence |
| return func |
|
|
| @staticmethod |
| def gen_forward(rewards, label_smoothing=0.0, top_k=1): |
| def func(self, x: Tensor, targets=None) -> Tensor: |
| x = self.conv1(x) |
| x = self.bn1(x) |
| x = self.relu(x) |
| x = self.maxpool(x) |
|
|
| x = self.layer1(x) |
| x = self.layer2(x) |
| x = self.layer3(x) |
| hidden_states = self.layer4(x) |
| x = self.avgpool(hidden_states) |
| x = torch.flatten(x, 1) |
| logits = self.fc(x) |
|
|
| if self.training: |
| shared_head = ResNetConfig.gen_shared_head(self) |
| compute_logits = ResNetConfig.gen_logits(self, shared_head) |
| compute_mask = ResNetConfig.gen_mask(label_smoothing, top_k) |
| compute_loss = ResNetConfig.gen_criterion(label_smoothing) |
| |
| logits_sequence = compute_logits(hidden_states) |
| mask_sequence = compute_mask(logits_sequence, targets) |
| loss_sequence = compute_loss(logits_sequence, targets) |
| loss = compute_policy_loss(loss_sequence, mask_sequence, rewards) |
| |
| return logits, loss |
|
|
| return logits |
| |
| return func |
|
|
|
|
| class MobileNetV2Config(ResNetConfig): |
| @staticmethod |
| def gen_shared_head(self): |
| def func(hidden_states): |
| """ |
| Args: |
| hidden_states (Tensor): Hidden States tensor of shape [B, C, H, W]. |
| |
| Returns: |
| logits (Tensor): Logits tensor of shape [B, C]. |
| """ |
| x = nn.functional.adaptive_avg_pool2d(hidden_states, (1, 1)) |
| x = torch.flatten(x, 1) |
| logits = self.classifier(x) |
| return logits |
| return func |
| |
| @staticmethod |
| def gen_forward(rewards, label_smoothing=0.0, top_k=1): |
| def func(self, x: Tensor, targets=None) -> Tensor: |
| hidden_states = self.features(x) |
| |
| x = nn.functional.adaptive_avg_pool2d(hidden_states, (1, 1)) |
| x = torch.flatten(x, 1) |
| logits = self.classifier(x) |
|
|
| if self.training: |
| shared_head = MobileNetV2Config.gen_shared_head(self) |
| compute_logits = MobileNetV2Config.gen_logits(self, shared_head) |
| compute_mask = MobileNetV2Config.gen_mask(label_smoothing, top_k) |
| compute_loss = MobileNetV2Config.gen_criterion(label_smoothing) |
| |
| logits_sequence = compute_logits(hidden_states) |
| mask_sequence = compute_mask(logits_sequence, targets) |
| loss_sequence = compute_loss(logits_sequence, targets) |
| loss = compute_policy_loss(loss_sequence, mask_sequence, rewards) |
| |
| return logits, loss |
|
|
| return logits |
| |
| return func |
|
|
|
|
| class EfficientNetConfig(ResNetConfig): |
| @staticmethod |
| def gen_shared_head(self): |
| def func(hidden_states): |
| """ |
| Args: |
| hidden_states (Tensor): Hidden States tensor of shape [B, C, H, W]. |
| |
| Returns: |
| logits (Tensor): Logits tensor of shape [B, C]. |
| """ |
| x = self.avgpool(hidden_states) |
| x = torch.flatten(x, 1) |
| logits = self.classifier(x) |
| return logits |
| return func |
| |
| @staticmethod |
| def gen_forward(rewards, label_smoothing=0.0, top_k=1): |
| def func(self, x: Tensor, targets=None) -> Tensor: |
| hidden_states = self.features(x) |
| x = self.avgpool(hidden_states) |
| x = torch.flatten(x, 1) |
| logits = self.classifier(x) |
|
|
| if self.training: |
| shared_head = EfficientNetConfig.gen_shared_head(self) |
| compute_logits = EfficientNetConfig.gen_logits(self, shared_head) |
| compute_mask = EfficientNetConfig.gen_mask(label_smoothing, top_k) |
| compute_loss = EfficientNetConfig.gen_criterion(label_smoothing) |
| |
| logits_sequence = compute_logits(hidden_states) |
| mask_sequence = compute_mask(logits_sequence, targets) |
| loss_sequence = compute_loss(logits_sequence, targets) |
| loss = compute_policy_loss(loss_sequence, mask_sequence, rewards) |
| |
| return logits, loss |
|
|
| return logits |
| |
| return func |
| |
|
|
| class VisionTransformerConfig(ResNetConfig): |
| @staticmethod |
| def gen_shared_head(self): |
| def func(hidden_states): |
| """ |
| Args: |
| hidden_states (Tensor): Hidden States tensor of shape [B, C, H, W]. |
| |
| Returns: |
| logits (Tensor): Logits tensor of shape [B, C]. |
| """ |
| x = hidden_states[:, 0] |
| logits = self.heads(x) |
| return logits |
| return func |
|
|
| @staticmethod |
| def gen_forward(rewards, label_smoothing=0.0, top_k=1): |
| def func(self, images: Tensor, targets=None): |
| x = self._process_input(images) |
| n = x.shape[0] |
| batch_class_token = self.class_token.expand(n, -1, -1) |
| x = torch.cat([batch_class_token, x], dim=1) |
| hidden_states = self.encoder(x) |
| x = hidden_states[:, 0] |
|
|
| logits = self.heads(x) |
|
|
|
|
| if self.training: |
| shared_head = VisionTransformerConfig.gen_shared_head(self) |
| compute_logits = VisionTransformerConfig.gen_logits(self, shared_head) |
| compute_mask = VisionTransformerConfig.gen_mask(label_smoothing, top_k) |
| compute_loss = VisionTransformerConfig.gen_criterion(label_smoothing) |
| |
| logits_sequence = compute_logits(hidden_states) |
| mask_sequence = compute_mask(logits_sequence, targets) |
| loss_sequence = compute_loss(logits_sequence, targets) |
| loss = compute_policy_loss(loss_sequence, mask_sequence, rewards) |
| |
| return logits, loss |
| return logits |
| return func |
| |
|
|
| def apply_trp(model, depths: List[int], in_planes: int, out_planes: int, rewards, **kwargs): |
| if isinstance(model, ResNet): |
| print("✅ Applying TRP to ResNet for Image Classification...") |
| model.trp_blocks = torch.nn.ModuleList([TPBlock(depths=d, in_planes=in_planes, out_planes=out_planes, rank=k) for k, d in enumerate(depths)]) |
| model.forward = types.MethodType(ResNetConfig.gen_forward(rewards, label_smoothing=kwargs["label_smoothing"], top_k=1), model) |
| elif isinstance(model, MobileNetV2): |
| print("✅ Applying TRP to MobileNetV2 for Image Classification...") |
| model.trp_blocks = torch.nn.ModuleList([TPBlock(depths=d, in_planes=in_planes, out_planes=out_planes, rank=k) for k, d in enumerate(depths)]) |
| model.forward = types.MethodType(MobileNetV2Config.gen_forward(rewards, label_smoothing=kwargs["label_smoothing"], top_k=1), model) |
| elif isinstance(model, EfficientNet): |
| print("✅ Applying TRP to EfficientNet for Image Classification...") |
| model.trp_blocks = torch.nn.ModuleList([TPBlock(depths=d, in_planes=in_planes, out_planes=out_planes, rank=k) for k, d in enumerate(depths)]) |
| model.forward = types.MethodType(EfficientNetConfig.gen_forward(rewards, label_smoothing=kwargs["label_smoothing"], top_k=1), model) |
| elif isinstance(model, VisionTransformer): |
| print("✅ Applying TRP to VisionTransformer for Image Classification...") |
| model.trp_blocks = torch.nn.ModuleList([TPBlock(depths=d, in_planes=in_planes, out_planes=out_planes, rank=k, shape_dims=2, channel_first=False) for k, d in enumerate(depths)]) |
| model.forward = types.MethodType(VisionTransformerConfig.gen_forward(rewards, label_smoothing=kwargs["label_smoothing"], top_k=1), model) |
| return model |
|
|
