initial push

example.py

@@ -0,0 +1,33 @@
+import torch
+from model.ResNet18 import ResNet18
+from model.CSAT import CSAT
+from model.CSATv2 import CSATv2
+from torch import nn
+
+img_size = 224
+path = r'./weight/CSAT_ImageNet.pth.tar' #or CSAT_RCKD.pth.tar <- for pathological image analysis
+model = CSAT(img_size=img_size)
+state = torch.load(path, map_location='cpu')
+model.load_state_dict(state)
+data = torch.zeros((1, 3, img_size, img_size)) #b, c, h, w = 1, 3, 224, 224
+model.head = nn.Identity()
+output = model(data)#b, c = 1, 176
+print(output.shape)
+
+path = r'./weight/ResNet18_RCKD.pth.tar'
+model = ResNet18()
+state = torch.load(path, map_location='cpu')
+model.load_state_dict(state)
+data = torch.zeros((1, 3, img_size, img_size)) #b, c, h, w = 1, 3, 224, 224
+model.fc = nn.Identity()
+output = model(data)#b, c = 1, 512
+print(output.shape)
+
+path = r'./weight/CSAT_v2_ImageNet.pth.tar'
+model = CSATv2(img_size=img_size)
+state = torch.load(path, map_location='cpu')
+model.load_state_dict(state['state_dict'])
+data = torch.zeros((1, 3, img_size, img_size)) #b, c, h, w = 1, 3, 224, 224
+model.fc = nn.Identity()
+output = model(data)#b, c = 1, 512
+print(output.shape)

model/CSAT.py

@@ -0,0 +1,490 @@
+import torch
+from torch import nn
+from einops.layers.torch import Rearrange
+from torch.nn.functional import softmax, sigmoid
+
+class Block(nn.Module):
+    """ ConvNeXtV2 Block.
+
+    Args:
+        dim (int): Number of input channels.
+        drop_path (float): Stochastic depth rate. Default: 0.0
+    """
+
+    def __init__(self, dim, drop_path=0., img_size=None):
+        super().__init__()
+        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)  # depthwise conv
+        self.norm = LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, 4 * dim)  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.grn = GRN(4 * dim)
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.attention = Spatial_Attention()
+    def forward(self, x):
+        input = x
+        x = self.dwconv(x)
+        x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pwconv2(x)
+
+        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)
+        attention = self.attention(x)
+        x = x * nn.UpsamplingBilinear2d(x.shape[2:])(attention)
+        x = input + self.drop_path(x)
+        return x
+
+class Spatial_Attention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.avgpool = nn.AdaptiveAvgPool2d((7,7))
+        self.conv = nn.Conv2d(2,1, kernel_size=7, padding=3)
+        self.attention = TransformerBlock(1, 1, heads=1, dim_head=1, img_size=[7,7])
+
+    def forward(self, x):
+        x_avg = x.mean([1]).unsqueeze(1)
+        x_max = x.max(dim=1).values.unsqueeze(1)
+        # x = torch.concat([x_avg,x_max],dim=1)
+        x = torch.cat([x_avg, x_max], dim=1)
+        x = self.avgpool(x)
+        x = self.conv(x)
+        x = self.attention(x)
+        return x
+
+class TransformerBlock(nn.Module):
+    def __init__(self, inp, oup, heads=8, dim_head=32, img_size=None, downsample=False, dropout=0.):
+        super().__init__()
+        hidden_dim = int(inp * 4)
+
+        self.downsample = downsample
+        self.ih, self.iw = img_size
+
+        if self.downsample:
+            self.pool1 = nn.MaxPool2d(3, 2, 1)
+            self.pool2 = nn.MaxPool2d(3, 2, 1)
+            self.proj = nn.Conv2d(inp, oup, 1, 1, 0, bias=False)
+
+        self.attn = Attention(inp, oup, heads, dim_head, dropout)
+        self.ff = FeedForward(oup, hidden_dim, dropout)
+
+        self.attn = nn.Sequential(
+            Rearrange('b c ih iw -> b (ih iw) c'),
+            PreNorm(inp, self.attn, nn.LayerNorm),
+            Rearrange('b (ih iw) c -> b c ih iw', ih=self.ih, iw=self.iw)
+        )
+
+        self.ff = nn.Sequential(
+            Rearrange('b c ih iw -> b (ih iw) c'),
+            PreNorm(oup, self.ff, nn.LayerNorm),
+            Rearrange('b (ih iw) c -> b c ih iw', ih=self.ih, iw=self.iw)
+        )
+
+    def forward(self, x):
+        if self.downsample:
+            x = self.proj(self.pool1(x)) + self.attn(self.pool2(x))
+        else:
+            x = x + self.attn(x)
+        x = x + self.ff(x)
+        return x
+
+
+class CSAT(nn.Module):
+    def __init__(self,
+                 img_size=384,
+                 num_classes=1000,
+                 drop_path_rate=0,
+                 head_init_scale=1,
+                 weight = None
+                 ):
+        super().__init__()
+        dims = [32, 48, 96, 176]
+        channel_order = "channels_first"
+        depths = [2, 2, 6, 4]
+        dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
+
+        self.stem = nn.Sequential(nn.Conv2d(in_channels=3, out_channels=dims[0], kernel_size=4, stride=4),
+                                  LayerNorm(normalized_shape=dims[0], data_format=channel_order))
+
+        self.stages1 = nn.Sequential(
+            Block(dim=dims[0], drop_path=dp_rates[0], img_size=[int(img_size / 4), int(img_size / 4)]),
+            Block(dim=dims[0], drop_path=dp_rates[1], img_size=[int(img_size / 4), int(img_size / 4)]),
+            LayerNorm(dims[0], eps=1e-6, data_format=channel_order),
+            nn.Conv2d(dims[0], dims[0 + 1], kernel_size=2, stride=2),
+        )
+
+        self.stages2 = nn.Sequential(
+            Block(dim=dims[1], drop_path=dp_rates[0], img_size=[int(img_size / 8), int(img_size / 8)]),
+            Block(dim=dims[1], drop_path=dp_rates[1], img_size=[int(img_size / 8), int(img_size / 8)]),
+            LayerNorm(dims[1], eps=1e-6, data_format=channel_order),
+            nn.Conv2d(dims[1], dims[1 + 1], kernel_size=2, stride=2),
+        )
+
+        self.stages3 = nn.Sequential(
+            Block(dim=dims[2], drop_path=dp_rates[0], img_size=[int(img_size / 16), int(img_size / 16)]),
+            Block(dim=dims[2], drop_path=dp_rates[1], img_size=[int(img_size / 16), int(img_size / 16)]),
+            Block(dim=dims[2], drop_path=dp_rates[2], img_size=[int(img_size / 16), int(img_size / 16)]),
+            Block(dim=dims[2], drop_path=dp_rates[3], img_size=[int(img_size / 16), int(img_size / 16)]),
+            Block(dim=dims[2], drop_path=dp_rates[4], img_size=[int(img_size / 16), int(img_size / 16)]),
+            Block(dim=dims[2], drop_path=dp_rates[5], img_size=[int(img_size / 16), int(img_size / 16)]),
+            TransformerBlock(inp=dims[2], oup=dims[2], img_size=[int(img_size / 16), int(img_size / 16)]),
+            TransformerBlock(inp=dims[2], oup=dims[2], img_size=[int(img_size / 16), int(img_size / 16)]),
+            LayerNorm(dims[2], eps=1e-6, data_format=channel_order),
+            nn.Conv2d(dims[2], dims[2 + 1], kernel_size=2, stride=2),
+        )
+
+        self.stages4 = nn.Sequential(
+            Block(dim=dims[3], drop_path=dp_rates[0], img_size=[int(img_size / 32), int(img_size / 32)]),
+            Block(dim=dims[3], drop_path=dp_rates[1], img_size=[int(img_size / 32), int(img_size / 32)]),
+            Block(dim=dims[3], drop_path=dp_rates[2], img_size=[int(img_size / 32), int(img_size / 32)]),
+            Block(dim=dims[3], drop_path=dp_rates[3], img_size=[int(img_size / 32), int(img_size / 32)]),
+            TransformerBlock(inp=dims[3], oup=dims[3], img_size=[int(img_size / 32), int(img_size / 32)]),
+            TransformerBlock(inp=dims[3], oup=dims[3], img_size=[int(img_size / 32), int(img_size / 32)]),
+        )
+
+        self.norm = nn.LayerNorm(dims[-1], eps=1e-6)  # final norm layer
+        self.head = nn.Linear(dims[-1], num_classes)
+
+        self.apply(self._init_weights)
+        self.head.weight.data.mul_(head_init_scale)
+        self.head.bias.data.mul_(head_init_scale)
+
+        if weight != None:
+            self.load_checkpoint(checkpoint=weight)
+        self.freeze_weight()
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Conv2d, nn.Linear)):
+            trunc_normal_(m.weight, std=.02)
+            try:
+                nn.init.constant_(m.bias, 0)
+            except:  # transformer layers
+                pass
+                # print("transformer layer can't initialize")
+
+    def freeze_weight(self):
+        for name, param in self.named_parameters():
+            if param.requires_grad and 'pos_embed' in name:
+                param.requires_grad = False
+
+    def load_checkpoint(self, checkpoint=None):
+        state = torch.load(checkpoint, map_location='cpu')
+        if 'state_dict' in state:
+            state_dict = state['state_dict']
+        elif 'model' in state:
+            state_dict = state['model']
+            for key in list(state_dict.keys()):
+                state_dict[key.replace('module.', '')] = state_dict.pop(key)
+        elif 'q_state_dict' in state:
+            state_dict = state['q_state_dict']
+
+        for key in list(state_dict.keys()):
+            state_dict[key.replace('backbone.', '')] = state_dict.pop(key)
+
+        model_dict = self.state_dict()
+        weights = {k: v for k, v in state_dict.items() if k in model_dict}
+
+        model_dict.update(weights)
+        del model_dict['head.weight']
+        del model_dict['head.bias']
+        self.load_state_dict(model_dict, strict=False)
+
+    def forward(self, x):
+        outputs = self.encoder(x)
+        # x, low_level, mid_level, high_level = self.seg_encoder(x)
+        return outputs
+
+    def encoder(self, x):
+        x = self.stem(x)
+        for _, layer in enumerate(self.stages1):
+            if _ == len(self.stages1) - 1:
+                x1 = x
+            x = layer(x)
+
+        for _, layer in enumerate(self.stages2):
+            if _ == len(self.stages2) - 1:
+                x2 = x
+            x = layer(x)
+
+        for _, layer in enumerate(self.stages3):
+            if _ == len(self.stages3) - 1:
+                x3 = x
+            x = layer(x)
+
+        x = self.stages4(x)
+        x = self.norm(x.mean([-2, -1]))
+        x = self.head(x)
+        return x
+
+    def seg_encoder(self, x):
+        org_img = x
+        x = self.stem(x)
+        for _, layer in enumerate(self.stages1):
+            if _ == len(self.stages1) - 2:
+                low_level = x
+            x = layer(x)
+
+        x = self.stages2(x)
+
+        for _, layer in enumerate(self.stages3):
+            if _ == len(self.stages3) - 2:
+                mid_level = x
+            x = layer(x)
+
+        for _, layer in enumerate(self.stages4):
+            x = layer(x)
+        high_level = x
+
+        return org_img, low_level, mid_level, high_level
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+import math
+import warnings
+
+class LayerNorm(nn.Module):
+    """ LayerNorm that supports two data formats: channels_last (default) or channels_first.
+    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
+    shape (batch_size, height, width, channels) while channels_first corresponds to inputs
+    with shape (batch_size, channels, height, width).
+    """
+
+    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        if self.data_format not in ["channels_last", "channels_first"]:
+            raise NotImplementedError
+        self.normalized_shape = (normalized_shape,)
+
+    def forward(self, x):
+        if self.data_format == "channels_last":
+            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        elif self.data_format == "channels_first":
+            u = x.mean(1, keepdim=True)
+            s = (x - u).pow(2).mean(1, keepdim=True)
+            x = (x - u) / torch.sqrt(s + self.eps)
+            x = self.weight[:, None, None] * x + self.bias[:, None, None]
+            return x
+
+
+class GRN(nn.Module):
+    """ GRN (Global Response Normalization) layer
+    """
+
+    def __init__(self, dim):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))
+
+    def forward(self, x):
+        Gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
+        Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
+        return self.gamma * (x * Nx) + self.beta + x
+
+def drop_path(x, drop_prob: float = 0., training: bool = False):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+
+    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
+    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
+    'survival rate' as the argument.
+
+    """
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()  # binarize
+    output = x.div(keep_prob) * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim, dropout=0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn, norm):
+        super().__init__()
+        self.norm = norm(dim)
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        return self.fn(self.norm(x), **kwargs)
+
+class Attention(nn.Module):
+    def __init__(self, inp, oup, heads=8, dim_head=32, dropout=0.):
+        super().__init__()
+        inner_dim = dim_head * heads
+        project_out = not (heads == 1 and dim_head == inp)
+
+        # self.ih, self.iw = image_size
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.attend = nn.Softmax(dim=-1)
+        self.to_qkv = nn.Linear(inp, inner_dim * 3, bias=False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, oup),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+        self.pos_embed = PosCNN(in_chans=inp)
+
+    def forward(self, x):
+        x = self.pos_embed(x)
+        qkv = self.to_qkv(x).chunk(3, dim=-1)
+        q, k, v = map(lambda t: rearrange(
+            t, 'b n (h d) -> b h n d', h=self.heads), qkv)
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+        attn = self.attend(dots)
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        out = self.to_out(out)
+        return out
+
+# PEG  from https://arxiv.org/abs/2102.10882
+class PosCNN(nn.Module):
+    def __init__(self, in_chans):
+        super(PosCNN, self).__init__()
+        self.proj = nn.Conv2d(in_chans, in_chans, kernel_size=3, stride = 1, padding=1, bias=True, groups=in_chans)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        feat_token = x
+        H, W = int(N**0.5), int(N**0.5)
+        cnn_feat = feat_token.transpose(1, 2).view(B, C, H, W)
+        x = self.proj(cnn_feat) + cnn_feat
+        x = x.flatten(2).transpose(1, 2)
+        return x
+
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    # type: (Tensor, float, float, float, float) -> Tensor
+    r"""Fills the input Tensor with values drawn from a truncated
+    normal distribution. The values are effectively drawn from the
+    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
+    with values outside :math:`[a, b]` redrawn until they are within
+    the bounds. The method used for generating the random values works
+    best when :math:`a \leq \text{mean} \leq b`.
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        mean: the mean of the normal distribution
+        std: the standard deviation of the normal distribution
+        a: the minimum cutoff value
+        b: the maximum cutoff value
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.trunc_normal_(w)
+    """
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+                      "The distribution of values may be incorrect.",
+                      stacklevel=2)
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        l = norm_cdf((a - mean) / std)
+        u = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [l, u], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+class DoubleConv(nn.Module):
+    """(convolution => [BN] => ReLU) * 2"""
+
+    def __init__(self, in_channels, out_channels, mid_channels=None):
+        super().__init__()
+        if not mid_channels:
+            mid_channels = out_channels
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(mid_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+
+class Up(nn.Module):
+    """Upscaling then double conv"""
+
+    def __init__(self, in_channels, out_channels, bilinear=True):
+        super().__init__()
+
+        # if bilinear, use the normal convolutions to reduce the number of channels
+        if bilinear:
+            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
+        else:
+            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
+            self.conv = DoubleConv(in_channels, out_channels)
+
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        # input is CHW
+        diffY = x2.size()[2] - x1.size()[2]
+        diffX = x2.size()[3] - x1.size()[3]
+
+        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
+                        diffY // 2, diffY - diffY // 2])
+        # if you have padding issues, see
+        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
+        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)

model/CSATv2.py

@@ -0,0 +1,396 @@
+import cv2
+import torch
+from torch import nn
+from einops.layers.torch import Rearrange
+from .DCT import Learnable_DCT2D #Learnable for H&E slide
+# from .DCT import Static_DCT2D #Static for Imagenet
+
+class Block(nn.Module):
+    """ ConvNeXtV2 Block.
+
+    Args:
+        dim (int): Number of input channels.
+        drop_path (float): Stochastic depth rate. Default: 0.0
+    """
+
+    def __init__(self, dim, drop_path=0.):
+        super().__init__()
+        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)  # depthwise conv
+        self.norm = LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, 4 * dim)  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.grn = GRN(4 * dim)
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.attention = Spatial_Attention()
+    def forward(self, x):
+        input = x
+        x = self.dwconv(x)
+        x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pwconv2(x)
+
+        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)
+        attention = self.attention(x)
+        x = x * nn.UpsamplingBilinear2d(x.shape[2:])(attention)
+        x = input + self.drop_path(x)
+        return x
+
+class Spatial_Attention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.avgpool = nn.AdaptiveAvgPool2d((7,7))
+        self.conv = nn.Conv2d(2,1, kernel_size=7, padding=3)
+        self.attention = TransformerBlock(1, 1, heads=1, dim_head=1, img_size=[7,7])
+
+    def forward(self, x):
+        x_avg = x.mean([1]).unsqueeze(1)
+        x_max = x.max(dim=1).values.unsqueeze(1)
+        # x = torch.concat([x_avg,x_max],dim=1)
+        x = torch.cat([x_avg, x_max], dim=1)
+        x = self.avgpool(x)
+        x = self.conv(x)
+        x = self.attention(x)
+        return x
+
+class TransformerBlock(nn.Module):
+    def __init__(self, inp, oup, heads=8, dim_head=32, img_size=None, downsample=False, dropout=0.):
+        super().__init__()
+        hidden_dim = int(inp * 4)
+
+        self.downsample = downsample
+        self.ih, self.iw = img_size
+
+        if self.downsample:
+            self.pool1 = nn.MaxPool2d(3, 2, 1)
+            self.pool2 = nn.MaxPool2d(3, 2, 1)
+            self.proj = nn.Conv2d(inp, oup, 1, 1, 0, bias=False)
+
+        self.attn = Attention(inp, oup, heads, dim_head, dropout)
+        self.ff = FeedForward(oup, hidden_dim, dropout)
+
+        self.attn = nn.Sequential(
+            Rearrange('b c ih iw -> b (ih iw) c'),
+            PreNorm(inp, self.attn, nn.LayerNorm),
+            Rearrange('b (ih iw) c -> b c ih iw', ih=self.ih, iw=self.iw)
+        )
+
+        self.ff = nn.Sequential(
+            Rearrange('b c ih iw -> b (ih iw) c'),
+            PreNorm(oup, self.ff, nn.LayerNorm),
+            Rearrange('b (ih iw) c -> b c ih iw', ih=self.ih, iw=self.iw)
+        )
+
+    def forward(self, x):
+        if self.downsample:
+            x = self.proj(self.pool1(x)) + self.attn(self.pool2(x))
+        else:
+            x = x + self.attn(x)
+        x = x + self.ff(x)
+        return x
+
+
+class CSATv2(nn.Module):
+    def __init__(self, img_size=None, num_classes=1000, drop_path_rate=0, head_init_scale=1):
+        super().__init__()
+        dims = [32, 72, 168, 386]
+        channel_order = "channels_first"
+        depths = [2, 2, 6, 4]
+        dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
+
+        # self.stem = nn.Sequential(nn.Conv2d(in_channels=3, out_channels=dims[0], kernel_size=4, stride=4),
+        #                           LayerNorm(normalized_shape=dims[0], data_format=channel_order))
+
+        self.stages1 = nn.Sequential(
+            Block(dim=dims[0], drop_path=dp_rates[0]),
+            Block(dim=dims[0], drop_path=dp_rates[1]),
+            LayerNorm(dims[0], eps=1e-6, data_format=channel_order),
+            nn.Conv2d(dims[0], dims[0 + 1], kernel_size=2, stride=2),
+        )
+
+        self.stages2 = nn.Sequential(
+            Block(dim=dims[1], drop_path=dp_rates[0]),
+            Block(dim=dims[1], drop_path=dp_rates[1]),
+            LayerNorm(dims[1], eps=1e-6, data_format=channel_order),
+            nn.Conv2d(dims[1], dims[1 + 1], kernel_size=2, stride=2),
+        )
+
+        self.stages3 = nn.Sequential(
+            Block(dim=dims[2], drop_path=dp_rates[0]),
+            Block(dim=dims[2], drop_path=dp_rates[1]),
+            Block(dim=dims[2], drop_path=dp_rates[2]),
+            Block(dim=dims[2], drop_path=dp_rates[3]),
+            Block(dim=dims[2], drop_path=dp_rates[4]),
+            Block(dim=dims[2], drop_path=dp_rates[5]),
+            TransformerBlock(inp=dims[2], oup=dims[2], img_size=[int(img_size / 32), int(img_size / 32)]),
+            TransformerBlock(inp=dims[2], oup=dims[2], img_size=[int(img_size / 32), int(img_size / 32)]),
+            LayerNorm(dims[2], eps=1e-6, data_format=channel_order),
+            nn.Conv2d(dims[2], dims[2 + 1], kernel_size=2, stride=2),
+        )
+
+        self.stages4 = nn.Sequential(
+            Block(dim=dims[3], drop_path=dp_rates[0]),
+            Block(dim=dims[3], drop_path=dp_rates[1]),
+            Block(dim=dims[3], drop_path=dp_rates[2]),
+            Block(dim=dims[3], drop_path=dp_rates[3]),
+            TransformerBlock(inp=dims[3], oup=dims[3], img_size=[int(img_size / 64), int(img_size / 64)]),
+            TransformerBlock(inp=dims[3], oup=dims[3], img_size=[int(img_size / 64), int(img_size / 64)]),
+        )
+
+        self.norm = nn.LayerNorm(dims[-1], eps=1e-6)  # final norm layer
+        self.head = nn.Linear(dims[-1], num_classes)
+
+        self.apply(self._init_weights)
+        self.head.weight.data.mul_(head_init_scale)
+        self.head.bias.data.mul_(head_init_scale)
+        self.dct = Learnable_DCT2D(8)
+        # self.dct = Static_DCT2D(8)
+
+    def load_checkpoint(self, checkpoint):
+        state = torch.load(checkpoint, map_location='cpu')
+        try:
+            state_dict = state['state_dict']
+        except:
+            state_dict = state['model']
+        for key in list(state_dict.keys()):
+            state_dict[key.replace('module.backbone.', '').replace('resnet.', '')] = state_dict.pop(key)
+
+        model_dict = self.state_dict()
+        weights = {k: v for k, v in state_dict.items() if k in model_dict}
+
+        model_dict.update(weights)
+        del model_dict['head.bias']
+        del model_dict['head.weight']
+        self.load_state_dict(model_dict, strict=False)
+
+    def preprocess(self,  x):
+        x = cv2.cvtColor(x, cv2.COLOR_BGR2YCR_CB)
+        return x
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Conv2d, nn.Linear)):
+            trunc_normal_(m.weight, std=.02)
+            try:
+                nn.init.constant_(m.bias, 0)
+            except:  # transformer layers
+                pass
+                # print("transformer layer can't initialize")
+
+
+    def forward(self, x):
+        # x = self.preprocess(x)
+        x = self.dct(x)#b, c, h, w -> b, c, *, h, w
+        x = self.stages1(x)
+        x = self.stages2(x)
+        x = self.stages3(x)
+        x = self.stages4(x)
+        x = self.norm(x.mean([-2, -1]))
+        x = self.head(x)
+        return x
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+import math
+import warnings
+
+class LayerNorm(nn.Module):
+    """ LayerNorm that supports two data formats: channels_last (default) or channels_first.
+    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
+    shape (batch_size, height, width, channels) while channels_first corresponds to inputs
+    with shape (batch_size, channels, height, width).
+    """
+
+    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        if self.data_format not in ["channels_last", "channels_first"]:
+            raise NotImplementedError
+        self.normalized_shape = (normalized_shape,)
+
+    def forward(self, x):
+        if self.data_format == "channels_last":
+            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        elif self.data_format == "channels_first":
+            u = x.mean(1, keepdim=True)
+            s = (x - u).pow(2).mean(1, keepdim=True)
+            x = (x - u) / torch.sqrt(s + self.eps)
+            x = self.weight[:, None, None] * x + self.bias[:, None, None]
+            return x
+
+
+class GRN(nn.Module):
+    """ GRN (Global Response Normalization) layer
+    """
+
+    def __init__(self, dim):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))
+
+    def forward(self, x):
+        Gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
+        Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
+        return self.gamma * (x * Nx) + self.beta + x
+
+def drop_path(x, drop_prob: float = 0., training: bool = False):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+
+    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
+    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
+    'survival rate' as the argument.
+
+    """
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()  # binarize
+    output = x.div(keep_prob) * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim, dropout=0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn, norm):
+        super().__init__()
+        self.norm = norm(dim)
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        return self.fn(self.norm(x), **kwargs)
+
+class Attention(nn.Module):
+    def __init__(self, inp, oup, heads=8, dim_head=32, dropout=0.):
+        super().__init__()
+        inner_dim = dim_head * heads
+        project_out = not (heads == 1 and dim_head == inp)
+
+        # self.ih, self.iw = image_size
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.attend = nn.Softmax(dim=-1)
+        self.to_qkv = nn.Linear(inp, inner_dim * 3, bias=False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, oup),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+        self.pos_embed = PosCNN(in_chans=inp)
+
+    def forward(self, x):
+        x = self.pos_embed(x)
+        qkv = self.to_qkv(x).chunk(3, dim=-1)
+        q, k, v = map(lambda t: rearrange(
+            t, 'b n (h d) -> b h n d', h=self.heads), qkv)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+        attn = self.attend(dots)
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        out = self.to_out(out)
+        return out
+
+# PEG  from https://arxiv.org/abs/2102.10882
+class PosCNN(nn.Module):
+    def __init__(self, in_chans):
+        super(PosCNN, self).__init__()
+        self.proj = nn.Conv2d(in_chans, in_chans, kernel_size=3, stride = 1, padding=1, bias=True, groups=in_chans)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        feat_token = x
+        H, W = int(N**0.5), int(N**0.5)
+        cnn_feat = feat_token.transpose(1, 2).view(B, C, H, W)
+        x = self.proj(cnn_feat) + cnn_feat
+        x = x.flatten(2).transpose(1, 2)
+        return x
+
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    # type: (Tensor, float, float, float, float) -> Tensor
+    r"""Fills the input Tensor with values drawn from a truncated
+    normal distribution. The values are effectively drawn from the
+    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
+    with values outside :math:`[a, b]` redrawn until they are within
+    the bounds. The method used for generating the random values works
+    best when :math:`a \leq \text{mean} \leq b`.
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        mean: the mean of the normal distribution
+        std: the standard deviation of the normal distribution
+        a: the minimum cutoff value
+        b: the maximum cutoff value
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.trunc_normal_(w)
+    """
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+                      "The distribution of values may be incorrect.",
+                      stacklevel=2)
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        l = norm_cdf((a - mean) / std)
+        u = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [l, u], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor

model/DCT.py

@@ -0,0 +1,265 @@
+from typing import List
+
+import numpy as np
+import torch
+import torch.nn as nn
+import math
+
+__all__ = ['DCT2D']
+
+
+# Helper Functions
+mean=[[932.42657,-0.00260,0.33415,-0.02840,0.00003,-0.02792,-0.00183,0.00006,0.00032,0.03402,-0.00571,0.00020,0.00006,-0.00038,-0.00558,-0.00116,-0.00000,-0.00047,-0.00008,-0.00030,0.00942,0.00161,-0.00009,-0.00006,-0.00014,-0.00035,0.00001,-0.00220,0.00033,-0.00002,-0.00003,-0.00020,0.00007,-0.00000,0.00005,0.00293,-0.00004,0.00006,0.00019,0.00004,0.00006,-0.00015,-0.00002,0.00007,0.00010,-0.00004,0.00008,0.00000,0.00008,-0.00001,0.00015,0.00002,0.00007,0.00003,0.00004,-0.00001,0.00004,-0.00000,0.00002,-0.00000,-0.00008,-0.00000,-0.00003,0.00003],
+[962.34735,-0.00428,0.09835,0.00152,-0.00009,0.00312,-0.00141,-0.00001,-0.00013,0.01050,0.00065,0.00006,-0.00000,0.00003,0.00264,0.00000,0.00001,0.00007,-0.00006,0.00003,0.00341,0.00163,0.00004,0.00003,-0.00001,0.00008,-0.00000,0.00090,0.00018,-0.00006,-0.00001,0.00007,-0.00003,-0.00001,0.00006,0.00084,-0.00000,-0.00001,0.00000,0.00004,-0.00001,-0.00002,0.00000,0.00001,0.00002,0.00001,0.00004,0.00011,0.00000,-0.00003,0.00011,-0.00002,0.00001,0.00001,0.00001,0.00001,-0.00007,-0.00003,0.00001,0.00000,0.00001,0.00002,0.00001,0.00000],
+[1053.16101,-0.00213,-0.09207,0.00186,0.00013,0.00034,-0.00119,0.00002,0.00011,-0.00984,0.00046,-0.00007,-0.00001,-0.00005,0.00180,0.00042,0.00002,-0.00010,0.00004,0.00003,-0.00301,0.00125,-0.00002,-0.00003,-0.00001,-0.00001,-0.00001,0.00056,0.00021,0.00001,-0.00001,0.00002,-0.00001,-0.00001,0.00005,-0.00070,-0.00002,-0.00002,0.00005,-0.00004,-0.00000,0.00002,-0.00002,0.00001,0.00000,-0.00003,0.00004,0.00007,0.00001,0.00000,0.00013,-0.00000,0.00000,0.00002,-0.00000,-0.00001,-0.00004,-0.00003,0.00000,0.00001,-0.00001,0.00001,-0.00000,0.00000]]
+
+var=[[270372.37500,6287.10645,5974.94043,1653.10889,1463.91748,1832.58997,755.92468,692.41528,648.57184,641.46881,285.79288,301.62100,380.43405,349.84027,374.15891,190.30960,190.76746,221.64578,200.82646,145.87979,126.92046,62.14622,67.75562,102.42001,129.74922,130.04631,103.12189,97.76417,53.17402,54.81048,73.48712,81.04342,69.35100,49.06024,33.96053,37.03279,20.48858,24.94830,33.90822,44.54912,47.56363,40.03160,30.43313,22.63899,26.53739,26.57114,21.84404,17.41557,15.18253,10.69678,11.24111,12.97229,15.08971,15.31646,8.90409,7.44213,6.66096,6.97719,4.17834,3.83882,4.51073,2.36646,2.41363,1.48266],
+[18839.21094,321.70932,300.15259,77.47830,76.02293,89.04748,33.99642,34.74807,32.12333,28.19588,12.04675,14.26871,18.45779,16.59588,15.67892,7.37718,8.56312,10.28946,9.41013,6.69090,5.16453,2.55186,3.03073,4.66765,5.85418,5.74644,4.33702,3.66948,1.95107,2.26034,3.06380,3.50705,3.06359,2.19284,1.54454,1.57860,0.97078,1.13941,1.48653,1.89996,1.95544,1.64950,1.24754,0.93677,1.09267,1.09516,0.94163,0.78966,0.72489,0.50841,0.50909,0.55664,0.63111,0.64125,0.38847,0.33378,0.30918,0.33463,0.20875,0.19298,0.21903,0.13380,0.13444,0.09554],
+[17127.39844,292.81421,271.45209,66.64056,63.60253,76.35437,28.06587,27.84831,25.96656,23.60370,9.99173,11.34992,14.46955,12.92553,12.69353,5.91537,6.60187,7.90891,7.32825,5.32785,4.29660,2.13459,2.44135,3.66021,4.50335,4.38959,3.34888,2.97181,1.60633,1.77010,2.35118,2.69018,2.38189,1.74596,1.26014,1.31684,0.79327,0.92046,1.17670,1.47609,1.50914,1.28725,0.99898,0.74832,0.85736,0.85800,0.74663,0.63508,0.58748,0.41098,0.41121,0.44663,0.50277,0.51519,0.31729,0.27336,0.25399,0.27241,0.17353,0.16255,0.18440,0.11602,0.11511,0.08450]]
+#torch.tensor(var)
+
+def _zigzag_permutation(rows: int, cols: int) -> List[int]:
+    idx_matrix = np.arange(0, rows * cols, 1).reshape(rows, cols).tolist()
+    dia = [[] for _ in range(rows + cols - 1)]
+    zigzag = []
+    for i in range(rows):
+        for j in range(cols):
+            s = i + j
+            if s % 2 == 0:
+                dia[s].insert(0, idx_matrix[i][j])
+            else:
+                dia[s].append(idx_matrix[i][j])
+    for d in dia:
+        zigzag.extend(d)
+    return zigzag
+
+
+# Kernels
+
+
+def _dct_kernel_type_2(
+    kernel_size: int,
+    orthonormal: bool,
+    device=None,
+    dtype=None,
+) -> torch.Tensor:
+    factory_kwargs = dict(device=device, dtype=dtype)
+    x = torch.eye(kernel_size, **factory_kwargs)
+    v = x.clone().contiguous().view(-1, kernel_size)
+    v = torch.cat([v, v.flip([1])], dim=-1)
+    v = torch.fft.fft(v, dim=-1)[:, :kernel_size]
+    try:
+        k = torch.tensor(-1j, **factory_kwargs) * torch.pi * torch.arange(kernel_size, **factory_kwargs)[None, :]
+    except:
+        k = torch.tensor(-1j, **factory_kwargs) * math.pi * torch.arange(kernel_size, **factory_kwargs)[None, :]
+    k = torch.exp(k / (kernel_size * 2))
+    v = v * k
+    v = v.real
+    if orthonormal:
+        v[:, 0] = v[:, 0] * torch.sqrt(torch.tensor(1 / (kernel_size * 4), **factory_kwargs))
+        v[:, 1:] = v[:, 1:] * torch.sqrt(torch.tensor(1 / (kernel_size * 2), **factory_kwargs))
+    v = v.contiguous().view(*x.shape)
+    return v
+
+
+def _dct_kernel_type_3(
+    kernel_size: int,
+    orthonormal: bool,
+    device=None,
+    dtype=None,
+) -> torch.Tensor:
+    return torch.linalg.inv(_dct_kernel_type_2(kernel_size, orthonormal, device, dtype))
+
+
+# Modules
+
+
+class _DCT1D(nn.Module):
+
+    def __init__(self, kernel_size: int, kernel_type: int = 2, orthonormal: bool = True,
+                 device=None, dtype=None) -> None:
+        factory_kwargs = dict(device=device, dtype=dtype)
+        super(_DCT1D, self).__init__()
+        kernel = {'2': _dct_kernel_type_2, '3': _dct_kernel_type_3}
+        self.weights = nn.Parameter(kernel[f'{kernel_type}'](kernel_size, orthonormal, **factory_kwargs).T, False)
+        self.register_parameter('bias', None)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return nn.functional.linear(x, self.weights, self.bias)
+
+
+class _DCT2D(nn.Module):
+
+    def __init__(self, kernel_size: int, kernel_type: int = 2, orthonormal: bool = True,
+                 device=None, dtype=None) -> None:
+        factory_kwargs = dict(device=device, dtype=dtype)
+        super(_DCT2D, self).__init__()
+        self.transform = _DCT1D(kernel_size, kernel_type, orthonormal, **factory_kwargs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # [..., H, W] @ DCT_Kernel.T -> [..., W, H] @ DCT_Kernel.T -> [..., H, W]
+        return self.transform(self.transform(x).transpose(-1, -2)).transpose(-1, -2)
+
+
+# Discrete Cosine Transforms (DCT)
+
+
+class Learnable_DCT2D(nn.Module):
+    r"""Computes the two-dimensional block-wise discrete cosine transform of :attr:`input`.
+
+    Args:
+        kernel_size (int): Size of the coefficient kernel
+        kernel_type (int): Type of the DCT (see Notes). Default: 2
+        orthonormal (bool): A boolean makes the corresponding matrix of coefficients orthonormal. Default: True
+
+    """
+
+    def __init__(self, kernel_size: int, kernel_type: int = 2, orthonormal: bool = True,
+                 device=None, dtype=None) -> None:
+        factory_kwargs = dict(device=device, dtype=dtype)
+        super(Learnable_DCT2D, self).__init__()
+        self.k = kernel_size
+        self.unfold = nn.Unfold(kernel_size=(kernel_size, kernel_size), stride=(kernel_size, kernel_size))
+        self.transform = _DCT2D(kernel_size, kernel_type, orthonormal, **factory_kwargs)
+        self.permutation = _zigzag_permutation(kernel_size, kernel_size)
+        self.Y_Conv = nn.Conv2d(kernel_size**2, 24, kernel_size=1, padding=0)
+        self.Cb_Conv = nn.Conv2d(kernel_size**2, 4, kernel_size=1, padding=0)
+        self.Cr_Conv = nn.Conv2d(kernel_size**2, 4, kernel_size=1, padding=0)
+        self.mean = torch.tensor(mean, requires_grad=False)
+        self.var = torch.tensor(var, requires_grad=False)
+        self.imagenet_mean = torch.tensor([0.485, 0.456, 0.406], requires_grad=False)
+        self.imagenet_var = torch.tensor([0.229, 0.224, 0.225], requires_grad=False)
+    def denormalize(self, x):
+        x = x.multiply(self.imagenet_var.to(x.device)).add_(self.imagenet_mean.to(x.device)) * 255  # denormalize
+        return x
+
+    def rgb2ycbcr(self, x):
+        y = (x[:,:,:,0] *0.299) + (x[:,:,:,1]* 0.587) + (x[:,:,:,2] * 0.114) #rgb2ycbcr
+        cb = 0.564 * (x[:,:,:,2] - y) + 128
+        cr = 0.713 * (x[:,:,:,0] - y) + 128
+        x = torch.stack([y, cb, cr],dim=-1)
+        return x
+
+    def frequncy_normalize(self, x):
+        x[:, 0, ].sub_(self.mean.to(x.device)[0]).div_((self.var.to(x.device)[0]**0.5+1e-8))
+        x[:, 1, ].sub_(self.mean.to(x.device)[1]).div_((self.var.to(x.device)[1]**0.5+1e-8))
+        x[:, 2, ].sub_(self.mean.to(x.device)[2]).div_((self.var.to(x.device)[2]**0.5+1e-8))
+        return x
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        b, c, h, w = x.shape #b, c, h, w
+        x = x.permute(0, 2, 3, 1)#b, h,  w, c
+        x = self.denormalize(x)#b, h, w,c
+        x = self.rgb2ycbcr(x)#b, h,  w, c
+        x = x.permute(0, 3, 1, 2)#b, c, h, w
+        x = self.unfold(x).transpose(-1, -2)#b,c, h, w -> b, c*block, blocks
+        x = x.reshape(b, h // self.k, w // self.k, c, self.k, self.k)
+        x = self.transform(x)
+        x = x.reshape(-1, c, self.k * self.k)
+        x = x[:, :, self.permutation]
+        x = self.frequncy_normalize(x)
+        x = x.reshape(b, h // self.k, w // self.k, c, -1)#? b, block -> b, h, w, c, block
+        x = x.permute(0, 3, 4, 1, 2).contiguous()  # b, c, block, h, w
+        x_Y = self.Y_Conv(x[:, 0, ])
+        x_Cb = self.Cb_Conv(x[:, 1, ])
+        x_Cr = self.Cr_Conv(x[:, 2, ])
+        x = torch.cat([x_Y, x_Cb, x_Cr], axis=1)
+        return x
+
+class Static_DCT2D(nn.Module):
+    r"""Computes the two-dimensional block-wise discrete cosine transform of :attr:`input`.
+
+    Args:
+        kernel_size (int): Size of the coefficient kernel
+        kernel_type (int): Type of the DCT (see Notes). Default: 2
+        orthonormal (bool): A boolean makes the corresponding matrix of coefficients orthonormal. Default: True
+
+    """
+
+    def __init__(self, kernel_size: int, kernel_type: int = 2, orthonormal: bool = True,
+                 device=None, dtype=None) -> None:
+        factory_kwargs = dict(device=device, dtype=dtype)
+        super(Static_DCT2D, self).__init__()
+        self.k = kernel_size
+        self.unfold = nn.Unfold(kernel_size=(kernel_size, kernel_size), stride=(kernel_size, kernel_size))
+        self.transform = _DCT2D(kernel_size, kernel_type, orthonormal, **factory_kwargs)
+        self.permutation = _zigzag_permutation(kernel_size, kernel_size)
+        self.mean = torch.tensor(mean, requires_grad=False)
+        self.var = torch.tensor(var, requires_grad=False)
+        self.imagenet_mean = torch.tensor([0.485, 0.456, 0.406], requires_grad=False)
+        self.imagenet_var = torch.tensor([0.229, 0.224, 0.225], requires_grad=False)
+
+    def denormalize(self, x):
+        x = x.multiply(self.imagenet_var.to(x.device)).add_(self.imagenet_mean.to(x.device)) * 255  # denormalize
+        return x
+
+    def rgb2ycbcr(self, x):
+        y = (x[:,:,:,0] *0.299) + (x[:,:,:,1]* 0.587) + (x[:,:,:,2] * 0.114) #rgb2ycbcr
+        cb = 0.564 * (x[:,:,:,2] - y) + 128
+        cr = 0.713 * (x[:,:,:,0] - y) + 128
+        x = torch.stack([y, cb, cr],dim=-1)
+        return x
+
+    def frequncy_normalize(self, x):
+        x[:, 0, ].sub_(self.mean.to(x.device)[0]).div_((self.var.to(x.device)[0]**0.5+1e-8))
+        x[:, 1, ].sub_(self.mean.to(x.device)[1]).div_((self.var.to(x.device)[1]**0.5+1e-8))
+        x[:, 2, ].sub_(self.mean.to(x.device)[2]).div_((self.var.to(x.device)[2]**0.5+1e-8))
+        return x
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        b, c, h, w = x.shape #b, c, h, w
+        x = x.permute(0, 2, 3, 1)#b, h,  w, c
+        x = self.denormalize(x)#b, h, w,c
+        x = self.rgb2ycbcr(x)#b, h,  w, c
+        x = x.permute(0, 3, 1, 2)#b, c, h, w
+        x = self.unfold(x).transpose(-1, -2)#b,c, h, w -> b, c*block, blocks
+        x = x.reshape(b, h // self.k, w // self.k, c, self.k, self.k)
+        x = self.transform(x)
+        x = x.reshape(-1, c, self.k * self.k)
+        x = x[:, :, self.permutation]
+        x = self.frequncy_normalize(x)
+        x = x.reshape(b, h // self.k, w // self.k, c, -1)#? b, block -> b, h, w, c, block
+        x = x.permute(0, 3, 4, 1, 2).contiguous()  # b, c, block, h, w
+        x_Y = self.Y_Conv(x[:, 0, ])
+        x_Cb = self.Cb_Conv(x[:, 1, ])
+        x_Cr = self.Cr_Conv(x[:, 2, ])
+        x = torch.cat([x_Y, x_Cb, x_Cr], axis=1)
+        return x
+
+class DCT2D(nn.Module):
+    r"""Computes the two-dimensional block-wise discrete cosine transform of :attr:`input`.
+
+    Args:
+        kernel_size (int): Size of the coefficient kernel
+        kernel_type (int): Type of the DCT (see Notes). Default: 2
+        orthonormal (bool): A boolean makes the corresponding matrix of coefficients orthonormal. Default: True
+
+    """
+
+    def __init__(self, kernel_size: int, kernel_type: int = 2, orthonormal: bool = True,
+                 device=None, dtype=None) -> None:
+        factory_kwargs = dict(device=device, dtype=dtype)
+        super(DCT2D, self).__init__()
+        self.k = kernel_size
+        self.unfold = nn.Unfold(kernel_size=(kernel_size, kernel_size), stride=(kernel_size, kernel_size))
+        self.transform = _DCT2D(kernel_size, kernel_type, orthonormal, **factory_kwargs)
+        self.permutation = _zigzag_permutation(kernel_size, kernel_size)
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        b, c, h, w = x.shape
+        x = self.unfold(x).transpose(-1, -2)#b,c, h, w -> b, c*block, blocks
+        x = x.reshape(b, h // self.k, w // self.k, c, self.k, self.k)
+        x = self.transform(x)
+        x = x.reshape(-1, c, self.k * self.k)
+        x = x[:, :, self.permutation]
+        x = x.reshape(b*(h // self.k)*(w // self.k), c, -1)#? b, block -> b, h, w, c, block
+        #torch.max(x[:,0,],axis=0).values.detach().cpu().numpy()
+
+        mean_list = torch.zeros([3,64])
+        var_list = torch.zeros([3, 64])
+        mean_list[0] = torch.mean(x[:, 0, ],axis=0)
+        mean_list[1] = torch.mean(x[:, 1, ], axis=0)
+        mean_list[2] = torch.mean(x[:, 2, ], axis=0)
+        var_list[0] = torch.var(x[:, 0, ],axis=0)
+        var_list[1] = torch.var(x[:, 1, ], axis=0)
+        var_list[2] = torch.var(x[:, 2, ], axis=0)
+        return mean_list, var_list

model/ResNet18.py

@@ -0,0 +1,9 @@
+import torchvision
+
+class ResNet18(torchvision.models.ResNet):
+    def __init__(self, num_classes=1000, weight=None):
+        super(ResNet18, self).__init__(block=torchvision.models.resnet.BasicBlock, layers=[2, 2, 2, 2], num_classes=num_classes)
+        self.zero_init_residual = True
+
+    def forward(self, x):
+        return self._forward_impl(x)

weight/CSAT_ImageNet.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f5ee778e954aa1a4f60a4a2f1376dbaa917bb42f9835d2170f2fd266f485d21c
+size 12417421

weight/CSAT_RCKD.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:efc69c12e2e11aec9487b06baf76cef4fd523bbdc34444b78240e19bb45337ce
+size 12417421

weight/CSAT_v2_ImageNet.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c34eb9183e3b4c89c0197ea870197b001313af420cd31f3f5304ed0e73a76e7
+size 44578564

weight/ResNet18_RCKD.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dfc3ca681cd4be87c0d4e4ed78669dd9b94e081eaab543d9e475666b7c1fee3b
+size 46836189