jacksung-dev 0.0.4.15__py3-none-any.whl

jacksung/ai/GeoNet/m_blockV2.py

@@ -0,0 +1,442 @@
+import math
+import random
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+import numpy as np
+
+
+class ShiftConv2d0(nn.Module):
+    def __init__(self, inp_channels, out_channels, stride):
+        super(ShiftConv2d0, self).__init__()
+        self.inp_channels = inp_channels
+        self.out_channels = out_channels
+        self.n_div = 5
+        self.stride = stride
+        g = inp_channels // self.n_div
+
+        conv3x3 = nn.Conv2d(inp_channels, out_channels, 3, 1, 1)
+        mask = nn.Parameter(torch.zeros((self.out_channels, self.inp_channels, 3, 3)), requires_grad=False)
+        mask[:, 0 * g:1 * g, 1, 2] = 1.0
+        mask[:, 1 * g:2 * g, 1, 0] = 1.0
+        mask[:, 2 * g:3 * g, 2, 1] = 1.0
+        mask[:, 3 * g:4 * g, 0, 1] = 1.0
+        mask[:, 4 * g:, 1, 1] = 1.0
+        self.w = conv3x3.weight
+        self.b = conv3x3.bias
+        self.m = mask
+
+    def forward(self, x):
+        y = F.conv2d(input=x, weight=self.w * self.m, bias=self.b, stride=self.stride, padding=1)
+        return y
+
+
+class ShiftConv2d1(nn.Module):
+    def __init__(self, inp_channels, out_channels, stride):
+        super(ShiftConv2d1, self).__init__()
+        self.inp_channels = inp_channels
+        self.out_channels = out_channels
+        self.stride = stride
+        self.weight = nn.Parameter(torch.zeros(inp_channels, 1, 3, 3), requires_grad=False)
+        self.n_div = 5
+        g = inp_channels // self.n_div
+
+        channels_idx = list(range(inp_channels))
+        random.shuffle(channels_idx)
+        self.weight[channels_idx[0 * g:1 * g], 0, 1, 2] = 1.0  ## left
+        self.weight[channels_idx[1 * g:2 * g], 0, 1, 0] = 1.0  ## right
+        self.weight[channels_idx[2 * g:3 * g], 0, 2, 1] = 1.0  ## up
+        self.weight[channels_idx[3 * g:4 * g], 0, 0, 1] = 1.0  ## down
+        self.weight[channels_idx[4 * g:], 0, 1, 1] = 1.0  ## identity
+        self.conv1x1 = nn.Conv2d(inp_channels, out_channels, 1)
+
+    def forward(self, x):
+        y = F.conv2d(input=x, weight=self.weight, bias=None, stride=self.stride, padding=1, groups=self.inp_channels)
+        y = self.conv1x1(y)
+        return y
+
+
+class ShiftConv2d(nn.Module):
+    def __init__(self, inp_channels, out_channels, conv_type='conv3', stride=1):
+        super(ShiftConv2d, self).__init__()
+        self.inp_channels = inp_channels
+        self.out_channels = out_channels
+        self.conv_type = conv_type
+        if conv_type == 'low-training-memory':
+            self.shift_conv = ShiftConv2d0(inp_channels, out_channels, stride=stride)
+        elif conv_type == 'fast-training-speed':
+            self.shift_conv = ShiftConv2d1(inp_channels, out_channels, stride=stride)
+        elif conv_type == 'common':
+            self.shift_conv = nn.Conv2d(inp_channels, out_channels, kernel_size=1, stride=stride)
+        elif conv_type == 'conv3':
+            self.shift_conv = nn.Conv2d(inp_channels, out_channels, kernel_size=3, stride=stride, padding=1)
+        else:
+            raise ValueError('invalid type of shift-conv2d')
+
+    def forward(self, x):
+        y = self.shift_conv(x)
+        return y
+
+
+class ACT(nn.Module):
+    def __init__(self):
+        super(ACT, self).__init__()
+        # self.act = nn.Mish()
+        self.act = nn.SiLU()
+
+    def forward(self, x):
+        return self.act(x)
+
+
+class DownBlock(nn.Module):
+    def __init__(self, c_lgan, downscale=2):
+        super(DownBlock, self).__init__()
+        self.down = nn.Conv2d(c_lgan, c_lgan * downscale, kernel_size=downscale, stride=downscale)
+        self.norm = Norm(c_lgan * downscale)
+        self.act = ACT()
+        self.conv = nn.Conv2d(c_lgan * downscale, c_lgan * downscale, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x):
+        x = self.down(x)
+        x = self.norm(x)
+        x = self.act(x)
+        x = self.conv(x)
+        return x
+
+
+class UpBlock(nn.Module):
+    def __init__(self, c_lgan, downscale=2):
+        super(UpBlock, self).__init__()
+        self.conv = nn.Conv2d(c_lgan * downscale, c_lgan * downscale, kernel_size=3, stride=1, padding=1)
+        self.norm = Norm(c_lgan * downscale)
+        self.act = ACT()
+        self.up = nn.ConvTranspose2d(c_lgan * downscale, c_lgan, kernel_size=downscale, stride=downscale)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.norm(x)
+        x = self.act(x)
+        x = self.up(x)
+        return x
+
+
+class Norm(nn.Module):
+    def __init__(self, c_in):
+        super(Norm, self).__init__()
+        self.norm = nn.BatchNorm2d(c_in)
+        # self.norm = nn.GroupNorm(1, c_in, eps=1e-6, affine=True)
+
+    def forward(self, x):
+        x = self.norm(x)
+        return x
+
+    # def __init__(self, c_in):
+    #     super(Norm, self).__init__()
+    #     self.norm = nn.LayerNorm(c_in, elementwise_affine=False)
+    #
+    # def forward(self, x):
+    #     b, c, h, w = x.shape
+    #     x = rearrange(x, 'b c h w->b (h w) c')
+    #     x = self.norm(x)
+    #     x = rearrange(x, 'b (h w) c->b c h w', w=w, h=h)
+    #     return x
+
+
+class CubeEmbeding(nn.Module):
+    def __init__(self, c_lgan, c_in, down_sample=1):
+        super(CubeEmbeding, self).__init__()
+        self.embeding = nn.Conv3d(c_in, c_lgan, kernel_size=(2, down_sample, down_sample),
+                                  stride=(2, down_sample, down_sample), padding=(0, 0, 0))
+        self.norm = Norm(c_lgan)
+        self.act = ACT()
+        self.conv2 = ShiftConv2d(c_lgan, c_lgan)
+
+    def forward(self, x):
+        x = self.embeding(x)
+        x = self.norm(x[:, :, 0, :, :])
+        x = self.act(x)
+        x = self.conv2(x)
+        return x
+
+
+class CubeUnEmbeding(nn.Module):
+    def __init__(self, c_lgan, c_in, downscale=2):
+        super(CubeUnEmbeding, self).__init__()
+        self.conv1 = nn.Conv2d(c_lgan, c_lgan, kernel_size=3, stride=1, padding=1)
+        self.norm = Norm(c_lgan)
+        self.act = ACT()
+        self.up = nn.ConvTranspose2d(c_lgan, c_lgan, kernel_size=downscale, stride=downscale)
+        self.conv2 = nn.Conv2d(c_lgan, c_in, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.norm(x)
+        x = self.act(x)
+        x = self.up(x)
+        x = self.conv2(x)
+        return x
+
+
+class Head(nn.Module):
+    def __init__(self, c_lgan, c_in, down_sample):
+        super(Head, self).__init__()
+        self.stage = nn.ModuleList()
+        while down_sample >= 2:
+            self.stage.append(
+                ShiftConv2d(c_in, c_in * 2, stride=2))
+            c_in *= 2
+            self.stage.append(Norm(c_in))
+            self.stage.append(ACT())
+            down_sample = down_sample // 2
+
+        self.conv2 = ShiftConv2d(c_in, c_lgan)
+
+    def forward(self, x):
+        # body
+        for stage in self.stage:
+            x = stage(x)
+        x = self.conv2(x)
+        return x
+
+
+class Tail(nn.Module):
+    def __init__(self, c_lgan, c_in, down_sample):
+        super(Tail, self).__init__()
+        self.down_sample = down_sample
+        self.conv1 = nn.Conv2d(c_lgan, c_lgan, kernel_size=3, stride=1, padding=1)
+        # self.norm = nn.BatchNorm2d(c_in * down_sample * down_sample)
+        self.norm = Norm(c_lgan)
+        self.act = ACT()
+        if self.down_sample == 1:
+            self.conv2 = nn.Conv2d(c_lgan, c_in, kernel_size=3,
+                                   stride=1, padding=1)
+        else:
+            self.conv2 = nn.Conv2d(c_lgan, c_in * down_sample * down_sample, kernel_size=3,
+                                   stride=1, padding=1)
+            self.ps = nn.PixelShuffle(down_sample)
+            # self.conv2 = nn.Conv2d(c_lgan, c_lgan * down_sample * down_sample, kernel_size=3,
+            #                        stride=1, padding=1)
+            # self.ps = nn.PixelShuffle(down_sample)
+            # self.conv3 = nn.Conv2d(c_lgan, c_in, kernel_size=3,
+            #                        stride=1, padding=1)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.norm(x)
+        x = self.act(x)
+        x = self.conv2(x)
+        if self.down_sample != 1:
+            x = self.ps(x)
+            # x = self.conv3(x)
+        return x
+
+
+# class Tail(nn.Module):
+#     def __init__(self, c_lgan, c_in, down_sample):
+#         super(Tail, self).__init__()
+#         self.c_in = c_in
+#         self.down_sample = down_sample
+#         self.fc = nn.Linear(c_lgan, c_in * down_sample * down_sample)
+#         self.norm = nn.LayerNorm(c_lgan)
+
+# def forward(self, x):
+#     b, c, h, w = x.shape
+#     x = rearrange(x, 'b c h w->b (h w) c')
+#     x = self.norm(x)
+#     x = self.fc(x)
+#     x = rearrange(x, 'b (h w) (c d1 d2)->b c (d1 h) (d2 w)',
+#                   c=self.c_in, d1=self.down_sample, d2=self.down_sample, h=h, w=w)
+#     x = nn.functional.interpolate(x, size=[721, 1440], mode='bilinear')
+#     return x
+
+
+class FD(nn.Module):
+    def __init__(self, inp_channels, out_channels, exp_ratio=4):
+        super(FD, self).__init__()
+        # self.fc1 = MLP(inp_channels, inp_channels * exp_ratio)
+        # self.fc2 = MLP(inp_channels * exp_ratio, out_channels)
+        self.fc1 = ShiftConv2d(inp_channels, inp_channels * exp_ratio)
+        self.fc2 = ShiftConv2d(inp_channels * exp_ratio, out_channels)
+        # self.fc1 = nn.Conv2d(inp_channels, inp_channels * exp_ratio, kernel_size=3, stride=1, padding=1)
+        # self.fc2 = nn.Conv2d(inp_channels * exp_ratio, out_channels, kernel_size=3, stride=1, padding=1)
+        self.act1 = ACT()
+        self.act2 = ACT()
+
+    def forward(self, x):
+        y = self.fc1(x)
+        y = self.act1(y)
+        y = self.fc2(y)
+        return y
+
+
+class LGAB(nn.Module):
+
+    def __init__(self, channels, window_size=5, num_heads=8, split_part=3):
+        super(LGAB, self).__init__()
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.split_chns = [int(channels * 2 / split_part) for _ in range(split_part)]
+        self.f_split_chns = [int(channels / split_part) for _ in range(split_part)]
+        # self.project_inp = MLP(channels, channels * 3)
+        # self.project_out = MLP(channels, channels)
+        self.project_inp = ShiftConv2d(channels, channels * 2)
+        self.f_project_inp = ShiftConv2d(channels, channels)
+        self.project_out = ShiftConv2d(channels, channels)
+        # self.f_project_inp = nn.Conv2d(channels, channels, kernel_size=1, stride=1)
+        # self.project_inp = nn.Conv2d(channels, channels * 2, kernel_size=1, stride=1)
+        # self.project_out = nn.Conv2d(channels, channels, kernel_size=1, stride=1)
+
+        self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((self.num_heads, 1, 1))), requires_grad=True)
+        self.lr_logit_scale = nn.Parameter(torch.log(10 * torch.ones((self.num_heads, 1, 1))), requires_grad=True)
+        # #########################################################################
+        # mlp to generate continuous relative position bias
+        self.cpb_mlp = nn.Sequential(nn.Linear(2, 512, bias=True),
+                                     nn.ReLU(inplace=True),
+                                     nn.Linear(512, self.num_heads, bias=False))
+
+        # get relative_coords_table
+        relative_coords_h = torch.arange(-(self.window_size - 1), self.window_size, dtype=torch.float32)
+        relative_coords_w = torch.arange(-(self.window_size - 1), self.window_size, dtype=torch.float32)
+        relative_coords_table = (torch.stack(
+            torch.meshgrid([relative_coords_h, relative_coords_w], indexing='ij'))
+                                 .permute(1, 2, 0).contiguous().unsqueeze(0))  # 1, 2*Wh-1, 2*Ww-1, 2
+
+        relative_coords_table[:, :, :, 0] /= (self.window_size - 1)
+        relative_coords_table[:, :, :, 1] /= (self.window_size - 1)
+        relative_coords_table *= 8  # normalize to -8, 8
+        relative_coords_table = torch.sign(relative_coords_table) * torch.log2(
+            torch.abs(relative_coords_table) + 1.0) / np.log2(8)
+
+        self.register_buffer("relative_coords_table", relative_coords_table)
+
+        # get pair-wise relative position index for each token inside the window
+        coords_h = torch.arange(self.window_size)
+        coords_w = torch.arange(self.window_size)
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing='ij'))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += self.window_size - 1  # shift to start from 0
+        relative_coords[:, :, 1] += self.window_size - 1
+        relative_coords[:, :, 0] *= 2 * self.window_size - 1
+        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+        self.register_buffer("relative_position_index", relative_position_index)
+        # #########################################################################
+
+    def wa(self, f, x, wsize):
+        b, c, h, w = x.shape
+        q = rearrange(
+            f, 'b (head c) (h dh) (w dw) -> (b h w) head (dh dw) c',
+            dh=wsize, dw=wsize, head=self.num_heads
+        )
+        k, v = rearrange(
+            x, 'b (kv head c) (h dh) (w dw) -> kv (b h w) head (dh dw) c',
+            kv=2, dh=wsize, dw=wsize, head=self.num_heads
+        )
+        atn = (F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1))
+        t = torch.tensor(1. / 0.01).to(atn.device)
+        logit_scale = torch.clamp(self.logit_scale, max=torch.log(t)).exp()
+        atn = atn * logit_scale
+
+        atn = atn.softmax(dim=-1)
+        y_ = (atn @ v)
+        y_ = rearrange(y_, '(b h w) head (dh dw) c-> b (head c) (h dh) (w dw)',
+                       h=h // wsize, w=w // wsize, dh=wsize, dw=wsize, head=self.num_heads)
+        return y_
+
+    # def forward(self, x, roll=False):
+    #     y_ = self.project_inp(x)
+    #     _, _, h, w = y_.shape
+    #     wsize = self.window_size
+    #     shifted window attention
+    # if roll:
+    #     y_ = torch.roll(y_, shifts=(-wsize // 2, -wsize // 2), dims=(2, 3))
+    # y_ = self.wa(y_, wsize, (h, w))
+    # if roll:
+    #     y_ = torch.roll(y_, shifts=(wsize // 2, wsize // 2), dims=(2, 3))
+    # y = self.project_out(y_)
+    # return y
+    def forward(self, x):
+        b, c, h, w = x.shape
+        x_ = x
+        x = self.project_inp(x_)
+        xs = torch.split(x, self.split_chns, dim=1)
+        f = self.f_project_inp(x_)
+        fs = torch.split(f, self.f_split_chns, dim=1)
+        wsize = self.window_size
+        ys = []
+        # window attention
+        y_ = self.wa(fs[0], xs[0], wsize)
+        ys.append(y_)
+        # shifted window attention
+        x_ = torch.roll(xs[1], shifts=(-wsize // 2, -wsize // 2), dims=(2, 3))
+        f_ = torch.roll(fs[1], shifts=(-wsize // 2, -wsize // 2), dims=(2, 3))
+        y_ = self.wa(f_, x_, wsize)
+        y_ = torch.roll(y_, shifts=(wsize // 2, wsize // 2), dims=(2, 3))
+        ys.append(y_)
+
+        # long-range attentin
+        # for longitude
+        q = rearrange(fs[2], 'b (head c) h w -> (b h) head w c', head=self.num_heads)
+        k, v = rearrange(xs[2], 'b (kv head c) h w -> kv (b h) head w c', kv=2, head=self.num_heads)
+        atn = (F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1))
+        t = torch.tensor(1. / 0.01).to(atn.device)
+        logit_scale = torch.clamp(self.lr_logit_scale, max=torch.log(t)).exp()
+        atn = atn * logit_scale
+        # atn = (q @ k.transpose(-2, -1))
+        atn = atn.softmax(dim=-1)
+        v = (atn @ v)
+        # for latitude
+        q, k, v = (rearrange(q, '(b h) head w c -> (b w) head h c', h=h),
+                   rearrange(k, '(b h) head w c -> (b w) head h c', h=h),
+                   rearrange(v, '(b h) head w c -> (b w) head h c', h=h))
+        atn = (F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1))
+        atn = atn * logit_scale
+        # atn = (q @ k.transpose(-2, -1))
+        atn = atn.softmax(dim=-1)
+        v = (atn @ v)
+        y_ = rearrange(v, '(b w) head h c-> b (head c) h w', b=b)
+        ys.append(y_)
+
+        y = torch.cat(ys, dim=1)
+        y = self.project_out(y)
+        return y
+
+
+class FEB(nn.Module):
+    def __init__(self, inp_channels, exp_ratio=2, window_size=5, num_heads=8):
+        super(FEB, self).__init__()
+        self.exp_ratio = exp_ratio
+        self.inp_channels = inp_channels
+        self.down = DownBlock(inp_channels, downscale=2)
+        self.up = UpBlock(inp_channels, downscale=2)
+
+        self.FD = FD(inp_channels=inp_channels * 2, out_channels=inp_channels * 2, exp_ratio=exp_ratio)
+        self.LGAB = LGAB(channels=inp_channels * 2, window_size=window_size, num_heads=num_heads)
+        self.norm1 = Norm(inp_channels * 2)
+        self.norm2 = Norm(inp_channels * 2)
+        self.drop = nn.Dropout2d(0.2)
+
+    def forward(self, x):
+        res = x
+        x = self.down(x)
+        shortcut = x
+        x = self.LGAB(x)
+        x = self.drop(x)
+        x = self.norm1(x) + shortcut
+        shortcut = x
+        x = self.FD(x)
+        x = self.norm2(x) + shortcut
+        x = self.up(x)
+        x = x + res
+        return x
+
+
+if __name__ == '__main__':
+    input_data = torch.zeros((5, 32, 32))
+    conv = nn.Conv2d(5, 5, stride=2, kernel_size=3, padding=1)
+    for i in range(5):
+        input_data = conv(input_data)
+        print(input_data.shape)

jacksung/ai/GeoNet/m_network.py

@@ -0,0 +1,107 @@
+import torch.nn as nn
+from jacksung.ai.GeoNet.m_block import FEB, Tail, Head, DownBlock, UpBlock, CubeEmbeding, CubeUnEmbeding, Norm, ACT
+import torch
+
+
+class GeoNet(nn.Module):
+    def __init__(self, window_sizes, n_lgab, c_in, c_lgan, r_expand=4, down_sample=4, num_heads=8, task='pred',
+                 downstage=2):
+        super(GeoNet, self).__init__()
+        self.window_sizes = window_sizes
+        self.n_lgab = n_lgab
+        self.c_in = c_in
+        self.c_lgan = c_lgan
+        self.r_expand = r_expand
+        self.task = task
+        self.down_sample = down_sample
+        # define head module
+        if self.task == 'prec':
+            self.head = Head(self.c_lgan, self.c_in, self.down_sample)
+        else:
+            self.head = CubeEmbeding(self.c_lgan, self.c_in, self.down_sample)
+            self.head_res = Head(self.c_lgan, self.c_in, self.down_sample)
+        # self.head = Head(self.c_lgan, self.c_in * 2, self.down_sample)
+        # define body module
+        self.body = nn.ModuleList()
+        self.downstage = downstage
+        for i in range(self.n_lgab):
+            if i / self.downstage in [1]:
+                self.body.append(DownBlock(self.c_lgan, 2))
+                self.c_lgan = self.c_lgan * 2
+            elif (self.n_lgab - i) / self.downstage in [1]:
+                self.body.append(UpBlock(self.c_lgan // 2, 2))
+                self.c_lgan = self.c_lgan // 2
+            self.body.append(
+                FEB(self.c_lgan, self.r_expand, self.window_sizes[i % len(self.window_sizes)], num_heads=num_heads))
+        # self.conv1 = nn.Conv2d(self.c_lgan, self.c_lgan, 3, 1, 1)
+        # self.norm = Norm(self.c_lgan)
+        # self.act = ACT()
+        # self.conv2 = nn.Conv2d(self.c_lgan, self.c_lgan, 3, 1, 1)
+        self.tail = Tail(self.c_lgan, 1 if self.task == 'prec' else self.c_in, down_sample)
+        # self.tail = nn.ConvTranspose3d(self.c_lgan, 5 if self.task == 'prec' else self.c_in, kernel_size=(2, 2, 2),
+        #                                stride=(2, 2, 2))
+        # self.tail = CubeUnEmbeding(self.c_lgan, 5 if self.task == 'prec' else self.c_in, self.down_sample)
+
+    def forward(self, f, x, roll=0):
+        # head
+        head_res = None
+        if self.task == 'pred':
+            head_res = self.head_res(x)
+            x = torch.stack([x, f], dim=2)
+        if roll > 0:
+            x = torch.roll(x, shifts=roll, dims=-1)
+        # f, x = rearrange(x, 'b c z h w->b (c z) h w'), rearrange(f, 'b c z h w->b (c z) h w')
+        # x = nn.functional.interpolate(x, size=[1600, 2000], mode='bilinear')
+        # f = nn.functional.interpolate(f, size=[1600, 2000], mode='bilinear')
+        x = self.head(x)
+        # shortcut = x
+        # body
+        x_res = None
+        for idx, stage in enumerate(self.body):
+            if idx / self.downstage in [1]:
+                x_res = x
+            x = stage(x)
+            if (self.n_lgab + 1 - idx) / self.downstage in [1]:
+                x += x_res
+        if self.task == 'pred':
+            x = head_res + x
+        # tail
+        x = self.tail(x)
+        # x = x[:, :, 0, :, :]
+        if roll > 0:
+            x = torch.roll(x, shifts=-roll, dims=-1)
+        # x = nn.functional.interpolate(x, size=[1607, 2008], mode='bilinear')
+        return x
+
+    def init_model(self):
+        print('Initializing the model!')
+        for m in self.children():
+            if isinstance(m, (nn.Conv2d, nn.Linear)):
+                nn.init.xavier_normal_(m.weight)
+
+    def load(self, state_dict, strict=True):
+        own_state = self.state_dict()
+        for name, param in state_dict.items():
+            name = name[name.index('.') + 1:]
+            if name in own_state.keys():
+                if isinstance(param, nn.Parameter):
+                    param = param.data
+                try:
+                    own_state[name].copy_(param)
+                    # own_state[name].requires_grad = False
+                except Exception as e:
+                    err_log = f'While copying the parameter named {name}, ' \
+                              f'whose dimensions in the model are {own_state[name].size()} and ' \
+                              f'whose dimensions in the checkpoint are {param.size()}.'
+                    if not strict:
+                        print(err_log)
+                    else:
+                        raise Exception(err_log)
+            elif strict:
+                raise KeyError(f'unexpected key {name} in {own_state.keys()}')
+            else:
+                print(f'{name} not loaded by model')
+
+
+if __name__ == '__main__':
+    pass

jacksung/ai/GeoNet/m_networkV2.py

@@ -0,0 +1,91 @@
+import torch.nn as nn
+from jacksung.ai.GeoNet.m_blockV2 import FEB, Tail, Head, DownBlock, UpBlock, CubeEmbeding, CubeUnEmbeding, Norm, ACT
+import torch
+import jacksung.utils.fastnumpy as fnp
+
+
+class GeoNet(nn.Module):
+    def __init__(self, window_sizes, n_lgab, c_in, c_lgan, r_expand=4, down_sample=2, num_heads=8, downstage=2):
+        super(GeoNet, self).__init__()
+        self.window_sizes = window_sizes
+        self.n_lgab = n_lgab
+        self.c_in = c_in
+        self.c_lgan = c_lgan
+        self.r_expand = r_expand
+        self.down_sample = down_sample
+        # define head module
+        self.head = Head(self.c_lgan, self.c_in, self.down_sample)
+        # define body module
+        self.body = nn.ModuleList()
+        self.downstage = downstage
+        for i in range(self.n_lgab):
+            if i % self.downstage in [1]:
+                if 0 <= i < self.n_lgab / 2 - 1:
+                    self.body.append(DownBlock(self.c_lgan, 2))
+                    self.c_lgan = self.c_lgan * 2
+                elif i > self.n_lgab / 2 + 1:
+                    self.body.append(UpBlock(self.c_lgan // 2, 2))
+                    self.c_lgan = self.c_lgan // 2
+            self.body.append(
+                FEB(self.c_lgan, self.r_expand, self.window_sizes[i % len(self.window_sizes)], num_heads=num_heads))
+        self.tail = Tail(self.c_lgan, 3, down_sample)
+
+    def forward(self, x, roll=0):
+        # head
+
+        x = torch.roll(x, shifts=roll, dims=-1)
+        x = self.head(x)
+        x_res = list()
+        x_res.append(x)
+        idx = 0
+        for stage in self.body:
+            if str(stage.__class__) == "<class 'jacksung.ai.GeoNet.m_blockV2.FEB'>":
+                # 7,9
+                if idx > self.n_lgab / 2 + 1 and idx % self.downstage in [1]:
+                    x += x_res.pop()
+                # 0,2
+                x = stage(x)
+                if idx < self.n_lgab / 2 - 1 and idx % self.downstage in [0]:
+                    x_res.append(x)
+                idx += 1
+            else:
+                x = stage(x)
+        # tail
+        x += x_res.pop()
+        x = self.tail(x)
+        if roll > 0:
+            x = torch.roll(x, shifts=-roll, dims=-1)
+        return x
+
+    def init_model(self):
+        print('Initializing the model!')
+        for m in self.children():
+            if isinstance(m, (nn.Conv2d, nn.Linear)):
+                nn.init.xavier_normal_(m.weight)
+
+    def load(self, state_dict, strict=True):
+        own_state = self.state_dict()
+        for name, param in state_dict.items():
+            name = name[name.index('.') + 1:]
+            if name in own_state.keys():
+                if isinstance(param, nn.Parameter):
+                    param = param.data
+                try:
+                    own_state[name].copy_(param)
+                    # own_state[name].requires_grad = False
+                except Exception as e:
+                    err_log = f'While copying the parameter named {name}, ' \
+                              f'whose dimensions in the model are {own_state[name].size()} and ' \
+                              f'whose dimensions in the checkpoint are {param.size()}.'
+                    if not strict:
+                        print(err_log)
+                    else:
+                        raise Exception(err_log)
+            elif strict:
+                raise KeyError(f'unexpected key {name} in {own_state.keys()}')
+            else:
+                print(f'{name} not loaded by model')
+
+
+if __name__ == '__main__':
+    pass