【25-Q4-生态建设】模型迁移-研发效能部-模型训练-在PyTorch框架上支持 TResNet 在Cifar100上的训练

PyTorch/build-in/Classification/TResNet/models/tresnet/__init__.py

@@ -0,0 +1 @@
+from .tresnet import TResnetM, TResnetL, TResnetXL

PyTorch/build-in/Classification/TResNet/models/tresnet/layers/anti_aliasing.py

@@ -0,0 +1,62 @@
+import torch
+import torch.nn.parallel
+import numpy as np
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class AntiAliasDownsampleLayer(nn.Module):
+    def __init__(self, remove_aa_jit: bool = False, filt_size: int = 3, stride: int = 2,
+                 channels: int = 0):
+        super(AntiAliasDownsampleLayer, self).__init__()
+        if not remove_aa_jit:
+            self.op = DownsampleJIT(filt_size, stride, channels)
+        else:
+            self.op = Downsample(filt_size, stride, channels)
+
+    def forward(self, x):
+        return self.op(x)
+
+
+@torch.jit.script
+class DownsampleJIT(object):
+    def __init__(self, filt_size: int = 3, stride: int = 2, channels: int = 0):
+        self.stride = stride
+        self.filt_size = filt_size
+        self.channels = channels
+
+        assert self.filt_size == 3
+        assert stride == 2
+        a = torch.tensor([1., 2., 1.])
+
+        filt = (a[:, None] * a[None, :]).clone().detach()
+        filt = filt / torch.sum(filt)
+        self.filt = filt[None, None, :, :].repeat((self.channels, 1, 1, 1)).cuda().half()
+
+    def __call__(self, input: torch.Tensor):
+        if input.dtype != self.filt.dtype:
+            self.filt = self.filt.float() 
+        input_pad = F.pad(input, (1, 1, 1, 1), 'reflect')
+        return F.conv2d(input_pad, self.filt, stride=2, padding=0, groups=input.shape[1])
+
+
+class Downsample(nn.Module):
+    def __init__(self, filt_size=3, stride=2, channels=None):
+        super(Downsample, self).__init__()
+        self.filt_size = filt_size
+        self.stride = stride
+        self.channels = channels
+
+
+        assert self.filt_size == 3
+        a = torch.tensor([1., 2., 1.])
+
+        filt = (a[:, None] * a[None, :])
+        filt = filt / torch.sum(filt)
+
+        # self.filt = filt[None, None, :, :].repeat((self.channels, 1, 1, 1))
+        self.register_buffer('filt', filt[None, None, :, :].repeat((self.channels, 1, 1, 1)))
+
+    def forward(self, input):
+        input_pad = F.pad(input, (1, 1, 1, 1), 'reflect')
+        return F.conv2d(input_pad, self.filt, stride=self.stride, padding=0, groups=input.shape[1])

PyTorch/build-in/Classification/TResNet/models/tresnet/layers/avg_pool.py

@@ -0,0 +1,19 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+
+class FastGlobalAvgPool2d(nn.Module):
+    def __init__(self, flatten=False):
+        super(FastGlobalAvgPool2d, self).__init__()
+        self.flatten = flatten
+
+    def forward(self, x):
+        if self.flatten:
+            in_size = x.size()
+            return x.view((in_size[0], in_size[1], -1)).mean(dim=2)
+        else:
+            return x.view(x.size(0), x.size(1), -1).mean(-1).view(x.size(0), x.size(1), 1, 1)
+
+

PyTorch/build-in/Classification/TResNet/models/tresnet/layers/space_to_depth.py

@@ -0,0 +1,53 @@
+import torch
+import torch.nn as nn
+
+
+class SpaceToDepth(nn.Module):
+    def __init__(self, block_size=4):
+        super().__init__()
+        assert block_size == 4
+        self.bs = block_size
+
+    def forward(self, x):
+        N, C, H, W = x.size()
+        x = x.view(N, C, H // self.bs, self.bs, W // self.bs, self.bs)  # (N, C, H//bs, bs, W//bs, bs)
+        x = x.permute(0, 3, 5, 1, 2, 4).contiguous()  # (N, bs, bs, C, H//bs, W//bs)
+        x = x.view(N, C * (self.bs ** 2), H // self.bs, W // self.bs)  # (N, C*bs^2, H//bs, W//bs)
+        return x
+
+
+@torch.jit.script
+class SpaceToDepthJit(object):
+    def __call__(self, x: torch.Tensor):
+        # assuming hard-coded that block_size==4 for acceleration
+        N, C, H, W = x.size()
+        x = x.view(N, C, H // 4, 4, W // 4, 4)  # (N, C, H//bs, bs, W//bs, bs)
+        x = x.permute(0, 3, 5, 1, 2, 4).contiguous()  # (N, bs, bs, C, H//bs, W//bs)
+        x = x.view(N, C * 16, H // 4, W // 4)  # (N, C*bs^2, H//bs, W//bs)
+        return x
+
+
+class SpaceToDepthModule(nn.Module):
+    def __init__(self, remove_model_jit=False):
+        super().__init__()
+        if not remove_model_jit:
+            self.op = SpaceToDepthJit()
+        else:
+            self.op = SpaceToDepth()
+
+    def forward(self, x):
+        return self.op(x)
+
+
+class DepthToSpace(nn.Module):
+
+    def __init__(self, block_size):
+        super().__init__()
+        self.bs = block_size
+
+    def forward(self, x):
+        N, C, H, W = x.size()
+        x = x.view(N, self.bs, self.bs, C // (self.bs ** 2), H, W)  # (N, bs, bs, C//bs^2, H, W)
+        x = x.permute(0, 3, 4, 1, 5, 2).contiguous()  # (N, C//bs^2, H, bs, W, bs)
+        x = x.view(N, C // (self.bs ** 2), H * self.bs, W * self.bs)  # (N, C//bs^2, H * bs, W * bs)
+        return x

PyTorch/build-in/Classification/TResNet/models/tresnet/layers/squeeze_and_excite.py

@@ -0,0 +1,39 @@
+import torch.nn as nn
+import torch.nn.functional as F
+from models.tresnet.layers.avg_pool import FastGlobalAvgPool2d
+
+
+class Flatten(nn.Module):
+    def forward(self, x):
+        return x.view(x.size(0), -1)
+
+
+class SEModule(nn.Module):
+
+    def __init__(self, channels, reduction_channels, inplace=True):
+        super(SEModule, self).__init__()
+        self.avg_pool = FastGlobalAvgPool2d()
+        self.fc1 = nn.Conv2d(channels, reduction_channels, kernel_size=1, padding=0, bias=True)
+        self.relu = nn.ReLU(inplace=inplace)
+        self.fc2 = nn.Conv2d(reduction_channels, channels, kernel_size=1, padding=0, bias=True)
+        # self.activation = hard_sigmoid(inplace=inplace)
+        self.activation = nn.Sigmoid()
+
+    def forward(self, x):
+        x_se = self.avg_pool(x)
+        x_se2 = self.fc1(x_se)
+        x_se2 = self.relu(x_se2)
+        x_se = self.fc2(x_se2)
+        x_se = self.activation(x_se)
+        return x * x_se
+
+class hard_sigmoid(nn.Module):
+    def __init__(self, inplace=True):
+        super(hard_sigmoid, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        if self.inplace:
+            return x.add_(3.).clamp_(0., 6.).div_(6.)
+        else:
+            return F.relu6(x + 3.) / 6.