| | 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 |
| |
|
