Upload model

AUNet.py

@@ -0,0 +1,212 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SqeezeExcite(nn.Module):
+    def __init__(self, channel, reduction_ratio = 16):
+        super(SqeezeExcite,self).__init__()
+        self.GAP = nn.AdaptiveAvgPool2d(1)
+
+        self.mlp = nn.Sequential(
+            nn.Linear(channel, channel//reduction_ratio, bias = False),
+            nn.ReLU(inplace=True),
+            nn.Linear(channel//reduction_ratio,channel,bias = False),
+            nn.Sigmoid()
+        )
+
+    def forward(self,x):
+        b,c,_,_ = x.size()
+        out = self.GAP(x).view(b,c)
+        out = self.mlp(out).view(b,c,1,1)
+        return x * out.expand_as(x)
+
+
+class ECA(nn.Module):
+    # https://wandb.ai/diganta/ECANet-sweep/reports/Efficient-Channel-Attention--VmlldzozNzgwOTE
+    def __init__(self,channels, b = 1, gamma = 2):
+        super(ECA, self).__init__()
+        self.GAP = nn.AdaptiveAvgPool2d(1)
+        self.channels = channels
+        self.b = b
+        self.gamma = gamma
+        self.conv = nn.Conv1d(1, 1, kernel_size=self.adaptive_kernel(),padding = (self.adaptive_kernel()-1)//2, bias = False)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self,x):
+        attn = self.GAP(x)
+        attn = self.conv(attn.squeeze(-1).transpose(-1,-2)).transpose(-1,-2).unsqueeze(-1)
+        attn = self.sigmoid(attn)
+        return x * attn.expand_as(x)
+
+
+    def adaptive_kernel(self):
+        k = int(abs(math.log2(self.channels)/self.gamma) + self.b)
+        ksize = k if k%2 else k+1
+        return ksize
+
+
+
+class UNetConvBlock(nn.Module):
+    def __init__(self, in_channel, out_channel, ca_layer):
+        super(UNetConvBlock, self).__init__()
+        block = []
+
+        block.append(nn.Conv2d(in_channel, out_channel, kernel_size=3, stride=1, padding=1))
+        block.append(nn.PReLU())
+
+        block.append(nn.Conv2d(out_channel, out_channel, kernel_size=3, padding=1, stride=1))
+        block.append(nn.PReLU())
+
+        if ca_layer:
+            block.append(ECA(out_channel))
+        self.block = nn.Sequential(*block)
+
+    def forward(self, x):
+        out = self.block(x)
+        return out
+
+class AttentionGate(nn.Module):
+    def __init__(self, F_g, F_l, dimensions):
+        super(AttentionGate, self).__init__()
+        self.W_gate = nn.Sequential(
+            nn.Conv2d(F_g, dimensions, kernel_size=1, stride=1, padding=0, bias=True),
+            nn.BatchNorm2d(dimensions)
+        )
+
+        self.W_x = nn.Sequential(
+            nn.Conv2d(F_l, dimensions, kernel_size=1, stride=1, padding=0, bias=True),
+            nn.BatchNorm2d(dimensions)
+        )
+
+        self.psi = nn.Sequential(
+            nn.Conv2d(dimensions, 1, kernel_size=1, stride=1, padding=0, bias=True),
+            nn.BatchNorm2d(1),
+            nn.Sigmoid()
+        )
+
+        self.relu = nn.PReLU()
+
+    def forward(self, g, x):
+        g1 = self.W_gate(g)
+        x1 = self.W_x(x)
+        psi = self.relu(g1 + x1)
+        psi = self.psi(psi)
+        out = x * psi
+        return out
+
+class UNetUpConvBlock(nn.Module):
+    def __init__(self, in_channel, out_channel, upmode, ca_layer, up_factor = 2, att_mode = "standard"):
+        super(UNetUpConvBlock, self).__init__()
+        self.att_mode = att_mode
+        self.ca_layer = ca_layer
+        if upmode == 'upsample':
+            self.Upsize = nn.Sequential(
+                                        nn.Upsample(scale_factor=up_factor, mode='bilinear', align_corners=False),
+                                        nn.Conv2d(in_channel, out_channel, kernel_size=1, stride=1, padding=0),
+                                      )
+        elif upmode == 'upconv':
+            self.Upsize = nn.ConvTranspose2d(in_channel,out_channel,kernel_size=2,stride = 2)
+        elif upmode == 'shuffle':
+            self.Upsize = nn.Sequential(
+                                        nn.Conv2d(in_channel,out_channel*4,kernel_size=3,stride=1,padding=1),
+                                        nn.PReLU(),
+                                        nn.PixelShuffle(2),
+                                        nn.Conv2d(out_channel,out_channel,kernel_size=3,stride = 1,padding=1)
+                                        )
+
+
+        # self.conv = UNetConvBlock(in_channel, out_channel)
+        if self.att_mode == 'standard':
+            self.attention_gate = AttentionGate(out_channel, out_channel, out_channel)
+            self.conv = UNetConvBlock(in_channel, out_channel, ca_layer=self.ca_layer)
+        elif self.att_mode == 'modified':
+            self.attention_gate = AttentionGate(out_channel, out_channel, out_channel )
+            self.conv = UNetConvBlock(3*out_channel, out_channel, ca_layer = self.ca_layer)
+        elif self.att_mode == 'None':
+            self.conv = UNetConvBlock(in_channel, out_channel, ca_layer=self.ca_layer)
+
+
+
+    def forward(self, x, residue):
+        x = self.Upsize(x)
+        x = F.interpolate(x, size=residue.shape[2:], mode='bilinear')
+        if self.att_mode == "standard":
+            attn = self.attention_gate(g = x, x=residue)
+            out = torch.cat([x, attn],dim = 1)
+            out = self.conv(out)
+        elif self.att_mode == 'modified':
+            attn = self.attention_gate(g = x, x = residue)
+            out = torch.cat([x,residue,attn],dim = 1)
+            out = self.conv(out)
+        elif self.att_mode == 'None':
+            out = torch.cat([x,residue], dim = 1)
+            out = self.conv(out)
+        return out
+
+
+
+
+
+class AUNet(nn.Module):
+    def __init__(self,in_channels = 6,out_channels = 6,depth = 3,growth_factor = 6,
+                 interp_mode = 'bicubic', up_mode = 'upconv',spatial_attention = "standard", ca_layer = True):
+        super(AUNet,self).__init__()
+
+        if not spatial_attention in ['None', 'modified', 'standard']:
+            raise AssertionError("spatial_attention options : \'None\'- no spatial attention, \'standard\'-spatial attention as in attention unet paper, \'modified\'-modified attention unet")
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.depth = depth
+        self.growth_factor = growth_factor
+        self.interp_mode = interp_mode
+        prev_channels = self.in_channels
+        self.up_mode = up_mode
+        self.att_mode = spatial_attention
+        self.ca_layer = ca_layer
+
+        self.encoding_module = nn.ModuleList()
+        for i in range(self.depth):
+            self.encoding_module.append(UNetConvBlock(in_channel=prev_channels,out_channel=2**(self.growth_factor + i), ca_layer=self.ca_layer))
+            prev_channels = 2**(self.growth_factor+i)
+
+        self.decoding_module = nn.ModuleList()
+        for i in reversed(range(self.depth-1)):
+            self.decoding_module.append(UNetUpConvBlock(prev_channels,2**(self.growth_factor+i),upmode = self.up_mode, att_mode = self.att_mode, ca_layer = self.ca_layer))
+            prev_channels = 2**(self.growth_factor+i)
+
+        self.final = nn.Conv2d(prev_channels,out_channels,1,1,0)
+
+    def forward(self,MS,PAN = None):
+        if PAN == None:
+            x = MS
+        else:
+            x = torch.cat([MS,PAN],dim = 1)
+        blocks = []
+        for i,down in enumerate(self.encoding_module):
+            x = down(x)
+            if i != len(self.encoding_module)-1:
+                blocks.append(x)
+                x = F.avg_pool2d(x,2)
+
+        for i,up in enumerate(self.decoding_module):
+            x = up(x,blocks[-i-1])
+
+        x = self.final(x)
+        return x
+
+
+if __name__ == '__main__':
+    x = torch.rand([9,7,256,256]).cuda()
+    model = AUNet(in_channels=7, out_channels=6, depth=5, spatial_attention="modified", growth_factor=6,
+                  interp_mode='bilinear', up_mode='upconv', ca_layer=True).cuda()
+    x = model(x)
+    # print(model)
+
+
+    activation = {}
+    for layer in model:
+        print(layer)

AUNetConfig.py

@@ -0,0 +1,28 @@
+from transformers import PretrainedConfig
+from typing import List
+
+
+class AUNetConfig(PretrainedConfig):
+    model_type = "s2l8hModel"
+    def __init__(
+            self,
+            in_channels:int = 7,
+            out_channels:int = 6,
+            depth:int = 5,
+            spatial_attention:str = 'None',
+            growth_factor:int = 6,
+            interp_mode:str = 'bicubic',
+            up_mode:str = 'upsample',
+            ca_layer:bool = False,
+            **kwargs,
+    ):
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.depth = depth
+        self.spatial_attention = spatial_attention
+        self.growth_factor = growth_factor
+        self.interp_mode = interp_mode
+        self.up_mode = up_mode
+        self.ca_layer = ca_layer
+
+        super().__init__(**kwargs)

config.json

@@ -0,0 +1,20 @@
+{
+  "architectures": [
+    "s2l8hModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "AUNetConfig.AUNetConfig",
+    "AutoModel": "s2l8hModel.s2l8hModel"
+  },
+  "ca_layer": false,
+  "depth": 6,
+  "growth_factor": 6,
+  "in_channels": 7,
+  "interp_mode": "bicubic",
+  "model_type": "s2l8hModel",
+  "out_channels": 6,
+  "spatial_attention": "None",
+  "torch_dtype": "float32",
+  "transformers_version": "4.33.1",
+  "up_mode": "upsample"
+}

pytorch_model.bin

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

s2l8hModel.py

@@ -0,0 +1,20 @@
+from transformers import PreTrainedModel
+from .AUNet import AUNet
+from .AUNetConfig import AUNetConfig
+import torch
+
+class s2l8hModel(PreTrainedModel):
+    config_class=AUNetConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = AUNet(
+            in_channels = config.in_channels, out_channels = config.out_channels,
+            depth = config.depth, spatial_attention = config.spatial_attention,
+            growth_factor = config.growth_factor, interp_mode = config.interp_mode,
+            up_mode = config.up_mode, ca_layer = config.ca_layer
+        )
+
+
+    def forward(self, MS, PAN):
+        return self.model.forward(MS, PAN)