Upload SCET.py

models/SCET.py

@@ -0,0 +1,276 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from einops import rearrange
+from einops.layers.torch import Rearrange
+import numbers
+
+# LayerNorm
+
+def to_3d(x):
+    return rearrange(x, 'b c h w -> b (h w) c')
+
+def to_4d(x,h,w):
+    return rearrange(x, 'b (h w) c -> b c h w',h=h,w=w)
+
+class BiasFree_LayerNorm(nn.Module):
+    def __init__(self, normalized_shape):
+        super(BiasFree_LayerNorm, self).__init__()
+        if isinstance(normalized_shape, numbers.Integral):
+            normalized_shape = (normalized_shape,)
+        normalized_shape = torch.Size(normalized_shape)
+
+        assert len(normalized_shape) == 1
+
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.normalized_shape = normalized_shape
+
+    def forward(self, x):
+        sigma = x.var(-1, keepdim=True, unbiased=False)
+        return x / torch.sqrt(sigma+1e-5) * self.weight
+
+class WithBias_LayerNorm(nn.Module):
+    def __init__(self, normalized_shape):
+        super(WithBias_LayerNorm, self).__init__()
+        if isinstance(normalized_shape, numbers.Integral):
+            normalized_shape = (normalized_shape,)
+        normalized_shape = torch.Size(normalized_shape)
+
+        assert len(normalized_shape) == 1
+
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.normalized_shape = normalized_shape
+
+    def forward(self, x):
+        mu = x.mean(-1, keepdim=True)
+        sigma = x.var(-1, keepdim=True, unbiased=False)
+        return (x - mu) / torch.sqrt(sigma+1e-5) * self.weight + self.bias
+
+class LayerNorm(nn.Module):
+    def __init__(self, dim, LayerNorm_type):
+        super(LayerNorm, self).__init__()
+        if LayerNorm_type =='BiasFree':
+            self.body = BiasFree_LayerNorm(dim)
+        else:
+            self.body = WithBias_LayerNorm(dim)
+
+    def forward(self, x):
+        h, w = x.shape[-2:]
+        return to_4d(self.body(to_3d(x)), h, w)
+
+
+## Gated-Dconv Feed-Forward Network (GDFN)
+class GFeedForward(nn.Module):
+    def __init__(self, dim, ffn_expansion_factor, bias):
+        super(GFeedForward, self).__init__()
+
+        hidden_features = int(dim * ffn_expansion_factor)
+
+        self.project_in = nn.Conv2d(dim, hidden_features * 2, kernel_size=1, bias=bias)
+
+        self.dwconv = nn.Conv2d(hidden_features * 2, hidden_features * 2, kernel_size=3, stride=1, padding=1,
+                                groups=hidden_features * 2, bias=bias)
+
+        self.project_out = nn.Conv2d(hidden_features, dim, kernel_size=1, bias=bias)
+
+    def forward(self, x):
+        x = self.project_in(x)
+        x1, x2 = self.dwconv(x).chunk(2, dim=1)
+        x = F.gelu(x1) * x2
+        x = self.project_out(x)
+        return x
+
+
+##########################################################################
+## Multi-DConv Head Transposed Self-Attention (MDTA)
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads, bias):
+        super(Attention, self).__init__()
+        self.num_heads = num_heads
+        self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))
+
+        self.qkv = nn.Conv2d(dim, dim * 3, kernel_size=1, bias=bias)
+        self.qkv_dwconv = nn.Conv2d(dim * 3, dim * 3, kernel_size=3, stride=1, padding=1, groups=dim * 3, bias=bias)
+        self.project_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias)
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+
+        qkv = self.qkv_dwconv(self.qkv(x))
+        q, k, v = qkv.chunk(3, dim=1)
+
+        q = rearrange(q, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+        k = rearrange(k, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+        v = rearrange(v, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+
+        q = torch.nn.functional.normalize(q, dim=-1)
+        k = torch.nn.functional.normalize(k, dim=-1)
+
+        attn = (q @ k.transpose(-2, -1)) * self.temperature
+        attn = attn.softmax(dim=-1)
+
+        out = (attn @ v)
+
+        out = rearrange(out, 'b head c (h w) -> b (head c) h w', head=self.num_heads, h=h, w=w)
+
+        out = self.project_out(out)
+        return out
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, dim=48, num_heads=8, ffn_expansion_factor=2.66, bias=False, LayerNorm_type=WithBias_LayerNorm):
+        super(TransformerBlock, self).__init__()
+
+        self.norm1 = LayerNorm(dim, LayerNorm_type)
+        self.attn = Attention(dim, num_heads, bias)
+        self.norm2 = LayerNorm(dim, LayerNorm_type)
+        self.ffn = GFeedForward(dim, ffn_expansion_factor, bias)
+
+    def forward(self, x):
+        x = x + self.attn(self.norm1(x))
+        x = x + self.ffn(self.norm2(x))
+
+        return x
+
+
+class BackBoneBlock(nn.Module):
+    def __init__(self, num, fm, **args):
+        super().__init__()
+        self.arr = nn.ModuleList([])
+        for _ in range(num):
+            self.arr.append(fm(**args))
+
+    def forward(self, x):
+        for block in self.arr:
+            x = block(x)
+        return x
+
+
+class PAConv(nn.Module):
+
+    def __init__(self, nf, k_size=3):
+        super(PAConv, self).__init__()
+        self.k2 = nn.Conv2d(nf, nf, 1)  # 1x1 convolution nf->nf
+        self.sigmoid = nn.Sigmoid()
+        self.k3 = nn.Conv2d(nf, nf, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)  # 3x3 convolution
+        self.k4 = nn.Conv2d(nf, nf, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)  # 3x3 convolution
+
+    def forward(self, x):
+        y = self.k2(x)
+        y = self.sigmoid(y)
+
+        out = torch.mul(self.k3(x), y)
+        out = self.k4(out)
+
+        return out
+
+
+class SCPA(nn.Module):
+    """SCPA is modified from SCNet (Jiang-Jiang Liu et al. Improving Convolutional Networks with Self-Calibrated Convolutions. In CVPR, 2020)
+        Github: https://github.com/MCG-NKU/SCNet
+    """
+
+    def __init__(self, nf, reduction=2, stride=1, dilation=1):
+        super(SCPA, self).__init__()
+        group_width = nf // reduction
+
+        self.conv1_a = nn.Conv2d(nf, group_width, kernel_size=1, bias=False)
+        self.conv1_b = nn.Conv2d(nf, group_width, kernel_size=1, bias=False)
+
+        self.k1 = nn.Sequential(
+            nn.Conv2d(
+                group_width, group_width, kernel_size=3, stride=stride,
+                padding=dilation, dilation=dilation,
+                bias=False)
+        )
+
+        self.PAConv = PAConv(group_width)
+
+        self.conv3 = nn.Conv2d(
+            group_width * reduction, nf, kernel_size=1, bias=False)
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+    def forward(self, x):
+        residual = x
+
+        out_a = self.conv1_a(x)
+        out_b = self.conv1_b(x)
+        out_a = self.lrelu(out_a)
+        out_b = self.lrelu(out_b)
+
+        out_a = self.k1(out_a)
+        out_b = self.PAConv(out_b)
+        out_a = self.lrelu(out_a)
+        out_b = self.lrelu(out_b)
+
+        out = self.conv3(torch.cat([out_a, out_b], dim=1))
+        out += residual
+
+        return out
+
+
+class SCET(nn.Module):
+    def __init__(self, hiddenDim=32, mlpDim=128, scaleFactor=2):
+        super().__init__()
+        self.conv3 = nn.Conv2d(3, hiddenDim,
+                               kernel_size=3, padding=1)
+
+        lamRes = torch.nn.Parameter(torch.ones(1))
+        lamX = torch.nn.Parameter(torch.ones(1))
+        self.adaptiveWeight = (lamRes, lamX)
+        if scaleFactor == 3:
+            num_heads = 7
+        else:
+            num_heads = 8
+        self.path1 = nn.Sequential(
+            BackBoneBlock(16, SCPA, nf=hiddenDim, reduction=2, stride=1, dilation=1),
+            BackBoneBlock(1, TransformerBlock,
+                          dim=hiddenDim, num_heads=num_heads, ffn_expansion_factor=2.66, bias=False, LayerNorm_type=WithBias_LayerNorm),
+            nn.Conv2d(hiddenDim, hiddenDim, kernel_size=3, padding=1),
+            nn.PixelShuffle(scaleFactor),
+            nn.Conv2d(hiddenDim // (scaleFactor ** 2),
+                      3, kernel_size=3, padding=1),
+        )
+
+        self.path2 = nn.Sequential(
+            nn.PixelShuffle(scaleFactor),
+            nn.Conv2d(hiddenDim // (scaleFactor ** 2),
+                      3, kernel_size=3, padding=1),
+        )
+
+    def forward(self, x):
+        x = self.conv3(x)
+        x1, x2 = self.path1(x), self.path2(x)
+        return x1 + x2
+
+
+    def init_weights(self, pretrained=None, strict=True):
+        """Init weights for models.
+        Args:
+            pretrained (str, optional): Path for pretrained weights. If given
+                None, pretrained weights will not be loaded. Defaults to None.
+            strict (boo, optional): Whether strictly load the pretrained model.
+                Defaults to True.
+        """
+        if isinstance(pretrained, str):
+            logger = get_root_logger()
+            load_checkpoint(self, pretrained, strict=strict, logger=logger)
+        elif pretrained is None:
+            pass  # use default initialization
+        else:
+            raise TypeError('"pretrained" must be a str or None. '
+                            f'But received {type(pretrained)}.')
+
+
+
+if __name__ == '__main__':
+
+    from torchstat import stat
+    import time
+    import torchsummary
+
+    net = SCET(32, 128, 4).cuda()
+    torchsummary.summary(net, (3, 48, 48))