Update submodules.py

submodules.py

@@ -1,833 +1,479 @@
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
-import torch.utils.checkpoint as checkpoint
-import numpy as np
-from timm.models.layers import DropPath, trunc_normal_
-
-from functools import reduce, lru_cache
-from operator import mul
-from einops import rearrange
-
-from model.submodules import ResidualBlock
-
-
-class residual_feature_generator(nn.Module):
-    def __init__(self, dim):
-        super(residual_feature_generator, self).__init__()
-        self.dim = dim
-        self.resblock1 = ResidualBlock(dim, dim, 1, norm='BN')
-        self.resblock2 = ResidualBlock(dim, dim, 1, norm='BN')
-        self.resblock3 = ResidualBlock(dim, dim, 1, norm='BN')
-        self.resblock4 = ResidualBlock(dim, dim, 1, norm='BN')
-    def forward(self, x):
-        out = self.resblock1(x)
-        out = self.resblock2(out)
-        out = self.resblock3(out)
-        out = self.resblock4(out)
-        return out
+import torch.nn.functional as f
+from torch.nn import init
+import math
+
+
+class ConvLayer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, activation='relu', norm=None,
+                 BN_momentum=0.1):
+        super(ConvLayer, self).__init__()
+
+        bias = False if norm == 'BN' else True
+        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias)
+        if activation is not None:
+            self.activation = getattr(torch, activation)
+        else:
+            self.activation = None
 
+        self.norm = norm
+        if norm == 'BN':
+            self.norm_layer = nn.BatchNorm2d(out_channels, momentum=BN_momentum)
+        elif norm == 'IN':
+            self.norm_layer = nn.InstanceNorm2d(out_channels, track_running_stats=True)
 
-class feature_generator(nn.Module):
-    def __init__(self, dim, kernel_size=3):
-        super(feature_generator, self).__init__()
-        self.dim = dim
-        self.kernel_size = kernel_size
-        self.conv1 = nn.Conv2d(in_channels=dim,
-                               out_channels=dim,
-                               kernel_size=kernel_size,
-                               stride=1,
-                               padding=(kernel_size-1)//2)
-        self.conv2 = nn.Conv2d(in_channels=dim,
-                               out_channels=dim,
-                               kernel_size=kernel_size,
-                               stride=1,
-                               padding=(kernel_size-1)//2)
-        self.conv3 = nn.Conv2d(in_channels=dim,
-                               out_channels=dim,
-                               kernel_size=kernel_size,
-                               stride=1,
-                               padding=(kernel_size-1)//2)
-        self.conv4 = nn.Conv2d(in_channels=dim,
-                               out_channels=dim,
-                               kernel_size=kernel_size,
-                               stride=1,
-                               padding=(kernel_size-1)//2)
-        self.bn1 = nn.BatchNorm2d(dim)
-        self.bn2 = nn.BatchNorm2d(dim)
-        self.bn3 = nn.BatchNorm2d(dim)
-        self.bn4 = nn.BatchNorm2d(dim)
     def forward(self, x):
-        out = F.leaky_relu(self.bn1(self.conv1(x)), negative_slope=0.01, inplace=False)
-        out = F.leaky_relu(self.bn2(self.conv2(out)), negative_slope=0.01, inplace=False)
-        out = F.leaky_relu(self.bn3(self.conv3(out)), negative_slope=0.01, inplace=False)
-        out = F.leaky_relu(self.bn4(self.conv4(out)), negative_slope=0.01, inplace=False)
+        out = self.conv2d(x)
+
+        if self.norm in ['BN', 'IN']:
+            out = self.norm_layer(out)
+
+        if self.activation is not None:
+            out = self.activation(out)
+
         return out
 
 
-class PatchEmbedLocalGlobal(nn.Module):
-    def __init__(self, patch_size=(2,4,4), in_chans=3, embed_dim=96, norm_layer=None):
-        super().__init__()
-        self.patch_size = patch_size
+class TransposedConvLayer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, activation='relu', norm=None):
+        super(TransposedConvLayer, self).__init__()
+
+        bias = False if norm == 'BN' else True
+        self.transposed_conv2d = nn.ConvTranspose2d(
+            in_channels, out_channels, kernel_size, stride=2, padding=padding, output_padding=1, bias=bias)
+
+        if activation is not None:
+            self.activation = getattr(torch, activation)
+        else:
+            self.activation = None
+
+        self.norm = norm
+        if norm == 'BN':
+            self.norm_layer = nn.BatchNorm2d(out_channels)
+        elif norm == 'IN':
+            self.norm_layer = nn.InstanceNorm2d(out_channels, track_running_stats=True)
+
+    def forward(self, x):
+        out = self.transposed_conv2d(x)
 
-        self.in_chans = in_chans
-        self.embed_dim = embed_dim
+        if self.norm in ['BN', 'IN']:
+            out = self.norm_layer(out)
 
-        self.num_blocks = self.in_chans // patch_size[0]
+        if self.activation is not None:
+            out = self.activation(out)
 
-        self.head = nn.Conv2d(in_chans // self.num_blocks, embed_dim // 2, kernel_size=3, stride=1, padding=1)
+        return out
 
-        self.global_head = nn.Conv2d(in_chans, embed_dim // 2, kernel_size=3, stride=1, padding=1)
 
-        self.residual_encoding = residual_feature_generator(embed_dim//2)
-        self.global_residual_encoding = residual_feature_generator(embed_dim//2)
+class UpsampleConvLayer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, activation='relu', norm=None):
+        super(UpsampleConvLayer, self).__init__()
 
-        self.proj = nn.Conv2d(embed_dim//2, embed_dim//2, kernel_size=3, stride=patch_size[1:], padding=1)
-        self.global_proj = nn.Conv2d(embed_dim//2, embed_dim//2, kernel_size=3, stride=patch_size[1:], padding=1)
+        bias = False if norm == 'BN' else True
+        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias)
 
-        if norm_layer is not None:
-            self.norm = norm_layer(embed_dim)
+        if activation is not None:
+            self.activation = getattr(torch, activation)
         else:
-            self.norm = None
+            self.activation = None
 
-        # patches_resolution = [224 // patch_size[1], 224 // patch_size[2]]
-        # self.absolute_pos_embed = nn.Parameter(torch.zeros(1, embed_dim, self.num_blocks, patches_resolution[0], patches_resolution[1]))
-        # trunc_normal_(self.absolute_pos_embed, std=.02)
+        self.norm = norm
+        if norm == 'BN':
+            self.norm_layer = nn.BatchNorm2d(out_channels)
+        elif norm == 'IN':
+            self.norm_layer = nn.InstanceNorm2d(out_channels, track_running_stats=True)
 
     def forward(self, x):
-        """Forward function."""
-        # padding
-        B, C, H, W = x.size()
-        # if W % self.patch_size[2] != 0:
-        #     x = F.pad(x, (0, self.patch_size[2] - W % self.patch_size[2]))
-        # if H % self.patch_size[1] != 0:
-        #     x = F.pad(x, (0, 0, 0, self.patch_size[1] - H % self.patch_size[1]))
-        # if D % self.patch_size[0] != 0:
-        #     x = F.pad(x, (0, 0, 0, 0, 0, self.patch_size[0] - D % self.patch_size[0]))
-        xs = x.chunk(self.num_blocks, 1)
-        outs = []
-        outi_global = self.global_head(x)
-        outi_global = self.global_residual_encoding(outi_global)
-        outi_global = self.global_proj(outi_global)
-
-        for i in range(self.num_blocks):
-            outi_local = self.head(xs[i])
-            outi_local = self.residual_encoding(outi_local)
-            outi_local = self.proj(outi_local)
-            outi = torch.cat([outi_local, outi_global], dim=1)
-            outi = outi.unsqueeze(2)
-            outs.append(outi)
-        
-        out = torch.cat(outs, dim=2) # B, 96, 4, H, W
-
-        # x = self.proj(x)  # B C D Wh Ww
-        if self.norm is not None:
-            D, Wh, Ww = out.size(2), out.size(3), out.size(4)
-            out = out.flatten(2).transpose(1, 2)
-            out = self.norm(out)
-            out = out.transpose(1, 2).view(-1, self.embed_dim, D, Wh, Ww)
+        x_upsampled = f.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)
+        out = self.conv2d(x_upsampled)
+
+        if self.norm in ['BN', 'IN']:
+            out = self.norm_layer(out)
+
+        if self.activation is not None:
+            out = self.activation(out)
 
         return out
 
 
-class PatchEmbedConv(nn.Module):
-    def __init__(self, patch_size=(2,4,4), in_chans=3, embed_dim=96, norm_layer=None):
-        super().__init__()
-        self.patch_size = patch_size
+class RecurrentConvLayer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=0,
+                 recurrent_block_type='convlstm', activation='relu', norm=None, BN_momentum=0.1):
+        super(RecurrentConvLayer, self).__init__()
+
+        assert(recurrent_block_type in ['convlstm', 'convgru'])
+        self.recurrent_block_type = recurrent_block_type
+        if self.recurrent_block_type == 'convlstm':
+            RecurrentBlock = ConvLSTM
+        else:
+            RecurrentBlock = ConvGRU
+
+        # self.conv = ConvLayer(in_channels, out_channels, kernel_size, stride, padding, activation, norm,
+        #                      BN_momentum=BN_momentum)
+        self.recurrent_block = RecurrentBlock(input_size=out_channels, hidden_size=out_channels, kernel_size=3)
+
+    def forward(self, x, prev_state):
+        # x = self.conv(x)
+        state = self.recurrent_block(x, prev_state)
+        x = state[0] if self.recurrent_block_type == 'convlstm' else state
+        return x, state
+
+class Recurrent2ConvLayer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=0,
+                 recurrent_block_type='convlstm', activation='relu', norm=None, BN_momentum=0.1):
+        super(Recurrent2ConvLayer, self).__init__()
+
+        assert(recurrent_block_type in ['convlstm', 'convgru'])
+        self.recurrent_block_type = recurrent_block_type
+        if self.recurrent_block_type == 'convlstm':
+            RecurrentBlock = ConvLSTM
+        else:
+            RecurrentBlock = ConvGRU
+
+        self.conv = ConvLayer(in_channels, out_channels, kernel_size, stride, padding, activation, norm,
+                              BN_momentum=BN_momentum)
+        self.recurrent_block = RecurrentBlock(input_size=out_channels, hidden_size=out_channels, kernel_size=3)
 
-        self.in_chans = in_chans
-        self.embed_dim = embed_dim
+    def forward(self, x, prev_state):
+        x = self.conv(x)
+        state = self.recurrent_block(x, prev_state)
+        x = state[0] if self.recurrent_block_type == 'convlstm' else state
+        return x, state
 
-        self.num_blocks = self.in_chans // patch_size[0]
 
-        self.head = nn.Conv2d(in_chans // self.num_blocks, embed_dim, kernel_size=3, stride=1, padding=1)
+class RecurrentPhasedConvLayer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=0,
+                 activation='relu', norm=None, BN_momentum=0.1):
+        super(RecurrentPhasedConvLayer, self).__init__()
 
-        self.residual_encoding = residual_feature_generator(embed_dim)
+        self.conv = ConvLayer(in_channels, out_channels, kernel_size, stride, padding, activation, norm,
+                              BN_momentum=BN_momentum)
+        self.recurrent_block = PhasedConvLSTMCell(input_channels=out_channels, hidden_channels=out_channels, kernel_size=3)
 
-        self.proj = nn.Conv2d(embed_dim, embed_dim, kernel_size=3, stride=patch_size[1:], padding=1)
-        if norm_layer is not None:
-            self.norm = norm_layer(embed_dim)
+    def forward(self, x, times, prev_state):
+        x = self.conv(x)
+        x, state = self.recurrent_block(x, times, prev_state)
+        return x, state
+
+
+class DownsampleRecurrentConvLayer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=3, recurrent_block_type='convlstm', padding=0, activation='relu'):
+        super(DownsampleRecurrentConvLayer, self).__init__()
+
+        self.activation = getattr(torch, activation)
+
+        assert(recurrent_block_type in ['convlstm', 'convgru'])
+        self.recurrent_block_type = recurrent_block_type
+        if self.recurrent_block_type == 'convlstm':
+            RecurrentBlock = ConvLSTM
         else:
-            self.norm = None
+            RecurrentBlock = ConvGRU
+        self.recurrent_block = RecurrentBlock(input_size=in_channels, hidden_size=out_channels, kernel_size=kernel_size)
+
+    def forward(self, x, prev_state):
+        state = self.recurrent_block(x, prev_state)
+        x = state[0] if self.recurrent_block_type == 'convlstm' else state
+        x = f.interpolate(x, scale_factor=0.5, mode='bilinear', align_corners=False)
+        return self.activation(x), state
+
+
+# Residual block
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride=1, downsample=None, norm=None,
+                 BN_momentum=0.1):
+        super(ResidualBlock, self).__init__()
+        bias = False if norm == 'BN' else True
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=bias)
+        self.norm = norm
+        if norm == 'BN':
+            self.bn1 = nn.BatchNorm2d(out_channels, momentum=BN_momentum)
+            self.bn2 = nn.BatchNorm2d(out_channels, momentum=BN_momentum)
+        elif norm == 'IN':
+            self.bn1 = nn.InstanceNorm2d(out_channels)
+            self.bn2 = nn.InstanceNorm2d(out_channels)
+
+        self.relu = nn.ReLU(inplace=False)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=bias)
+        self.downsample = downsample
 
     def forward(self, x):
-        """Forward function."""
-        # padding
-        B, C, H, W = x.size()
-        # if W % self.patch_size[2] != 0:
-        #     x = F.pad(x, (0, self.patch_size[2] - W % self.patch_size[2]))
-        # if H % self.patch_size[1] != 0:
-        #     x = F.pad(x, (0, 0, 0, self.patch_size[1] - H % self.patch_size[1]))
-        # if D % self.patch_size[0] != 0:
-        #     x = F.pad(x, (0, 0, 0, 0, 0, self.patch_size[0] - D % self.patch_size[0]))
-        xs = x.chunk(self.num_blocks, 1)
-        outs = []
-        
-        for i in range(self.num_blocks):
-            outi = self.head(xs[i])
-            outi = self.residual_encoding(outi)
-            outi = self.proj(outi)
-            outi = outi.unsqueeze(2)
-            outs.append(outi)
-        
-        out = torch.cat(outs, dim=2) # B, 96, 4, H, W
-
-        # x = self.proj(x)  # B C D Wh Ww
-        if self.norm is not None:
-            D, Wh, Ww = out.size(2), out.size(3), out.size(4)
-            out = out.flatten(2).transpose(1, 2)
-            out = self.norm(out)
-            out = out.transpose(1, 2).view(-1, self.embed_dim, D, Wh, Ww)
-
+        residual = x
+        out = self.conv1(x)
+        if self.norm in ['BN', 'IN']:
+            out = self.bn1(out)
+        out = self.relu(out)
+        out = self.conv2(out)
+        if self.norm in ['BN', 'IN']:
+            out = self.bn2(out)
+
+        if self.downsample:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
         return out
 
 
-class Mlp(nn.Module):
-    """ Multilayer perceptron."""
+class PhasedLSTMCell(nn.Module):
+    """Phased LSTM recurrent network cell.
+    """
 
-    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+    def __init__(
+        self,
+        hidden_size,
+        leak=0.001,
+        ratio_on=0.1,
+        period_init_min=0.02,
+        period_init_max=50.0
+    ):
+        """
+        Args:
+            hidden_size: int, The number of units in the Phased LSTM cell.
+            leak: float or scalar float Tensor with value in [0, 1]. Leak applied
+                during training.
+            ratio_on: float or scalar float Tensor with value in [0, 1]. Ratio of the
+                period during which the gates are open.
+            period_init_min: float or scalar float Tensor. With value > 0.
+                Minimum value of the initialized period.
+                The period values are initialized by drawing from the distribution:
+                e^U(log(period_init_min), log(period_init_max))
+                Where U(.,.) is the uniform distribution.
+            period_init_max: float or scalar float Tensor.
+                With value > period_init_min. Maximum value of the initialized period.
+        """
         super().__init__()
-        out_features = out_features or in_features
-        hidden_features = hidden_features or in_features
-        self.fc1 = nn.Linear(in_features, hidden_features)
-        self.act = act_layer()
-        self.fc2 = nn.Linear(hidden_features, out_features)
-        self.drop = nn.Dropout(drop)
 
-    def forward(self, x):
-        x = self.fc1(x)
-        x = self.act(x)
-        x = self.drop(x)
-        x = self.fc2(x)
-        x = self.drop(x)
-        return x
+        self.hidden_size = hidden_size
+        self.ratio_on = ratio_on
+        self.leak = leak
 
+        # initialize time-gating parameters
+        period = torch.exp(
+            torch.Tensor(hidden_size).uniform_(
+                math.log(period_init_min), math.log(period_init_max)
+            )
+        )
+        #self.tau = nn.Parameter(period)
+        self.register_parameter("tau", nn.Parameter(period))
 
-def window_partition(x, window_size):
-    """
-    Args:
-        x: (B, D, H, W, C)
-        window_size (tuple[int]): window size
-    Returns:
-        windows: (B*num_windows, window_size*window_size, C)
-    """
-    B, D, H, W, C = x.shape
-    x = x.view(B, D // window_size[0], window_size[0], H // window_size[1], window_size[1], W // window_size[2], window_size[2], C)
-    windows = x.permute(0, 1, 3, 5, 2, 4, 6, 7).contiguous().view(-1, reduce(mul, window_size), C)
-    return windows
+        phase = torch.Tensor(hidden_size).uniform_() * period
+        self.register_parameter("phase", nn.Parameter(phase))
+        self.phase.requires_grad = True
+        self.tau.requires_grad = True
+        #self.phase = nn.Parameter(phase)
 
+    def _compute_phi(self, t):
+        t_ = t.view(-1, 1).repeat(1, self.hidden_size)
+        phase_ = self.phase.view(1, -1).repeat(t.shape[0], 1)
+        tau_ = self.tau.view(1, -1).repeat(t.shape[0], 1)
+        tau_.to(t_.device)
+        phase_.to(t_.device)
+        phi = self._mod((t_ - phase_), tau_)
+        phi = torch.abs(phi) / tau_
+        return phi
 
-def window_reverse(windows, window_size, B, D, H, W):
-    """
-    Args:
-        windows: (B*num_windows, window_size, window_size, C)
-        window_size (tuple[int]): Window size
-        H (int): Height of image
-        W (int): Width of image
-    Returns:
-        x: (B, D, H, W, C)
-    """
-    x = windows.view(B, D // window_size[0], H // window_size[1], W // window_size[2], window_size[0], window_size[1], window_size[2], -1)
-    x = x.permute(0, 1, 4, 2, 5, 3, 6, 7).contiguous().view(B, D, H, W, -1)
-    return x
-
-
-def get_window_size(x_size, window_size, shift_size=None):
-    use_window_size = list(window_size)
-    if shift_size is not None:
-        use_shift_size = list(shift_size)
-    for i in range(len(x_size)):
-        if x_size[i] <= window_size[i]:
-            use_window_size[i] = x_size[i]
-            if shift_size is not None:
-                use_shift_size[i] = 0
-
-    if shift_size is None:
-        return tuple(use_window_size)
-    else:
-        return tuple(use_window_size), tuple(use_shift_size)
-
-
-class WindowAttention3D(nn.Module):
-    """ Window based multi-head self attention (W-MSA) module with relative position bias.
-    It supports both of shifted and non-shifted window.
-    Args:
-        dim (int): Number of input channels.
-        window_size (tuple[int]): The temporal length, height and width of the window.
-        num_heads (int): Number of attention heads.
-        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
-        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
-        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
-        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
-    """
+    def _mod(self, x, y):
+        """Modulo function that propagates x gradients."""
+        return x + (torch.fmod(x, y) - x).detach()
 
-    def __init__(self, dim, window_size, num_heads, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+    def set_state(self, c, h):
+        self.h0 = h
+        self.c0 = c
 
-        super().__init__()
-        self.dim = dim
-        self.window_size = window_size  # Wd, Wh, Ww
-        self.num_heads = num_heads
-        head_dim = dim // num_heads
-        self.scale = qk_scale or head_dim ** -0.5
-
-        # define a parameter table of relative position bias
-        self.relative_position_bias_table = nn.Parameter(
-            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1) * (2 * window_size[2] - 1), num_heads))  # 2*Wd-1 * 2*Wh-1 * 2*Ww-1, nH
-
-        # get pair-wise relative position index for each token inside the window
-        coords_d = torch.arange(self.window_size[0])
-        coords_h = torch.arange(self.window_size[1])
-        coords_w = torch.arange(self.window_size[2])
-        coords = torch.stack(torch.meshgrid(coords_d, coords_h, coords_w))  # 3, Wd, Wh, Ww
-        coords_flatten = torch.flatten(coords, 1)  # 3, Wd*Wh*Ww
-        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 3, Wd*Wh*Ww, Wd*Wh*Ww
-        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wd*Wh*Ww, Wd*Wh*Ww, 3
-        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
-        relative_coords[:, :, 1] += self.window_size[1] - 1
-        relative_coords[:, :, 2] += self.window_size[2] - 1
-
-        relative_coords[:, :, 0] *= (2 * self.window_size[1] - 1) * (2 * self.window_size[2] - 1)
-        relative_coords[:, :, 1] *= (2 * self.window_size[2] - 1)
-        relative_position_index = relative_coords.sum(-1)  # Wd*Wh*Ww, Wd*Wh*Ww
-        self.register_buffer("relative_position_index", relative_position_index)
-
-        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
-        self.attn_drop = nn.Dropout(attn_drop)
-        self.proj = nn.Linear(dim, dim)
-        self.proj_drop = nn.Dropout(proj_drop)
-
-        trunc_normal_(self.relative_position_bias_table, std=.02)
-        self.softmax = nn.Softmax(dim=-1)
-
-    def forward(self, x, mask=None):
-        """ Forward function.
-        Args:
-            x: input features with shape of (num_windows*B, N, C)
-            mask: (0/-inf) mask with shape of (num_windows, N, N) or None
-        """
-        B_, N, C = x.shape
-        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
-        q, k, v = qkv[0], qkv[1], qkv[2]  # B_, nH, N, C
-
-        q = q * self.scale
-        attn = q @ k.transpose(-2, -1)
-
-        relative_position_bias = self.relative_position_bias_table[self.relative_position_index[:N, :N].reshape(-1)].reshape(
-            N, N, -1)  # Wd*Wh*Ww,Wd*Wh*Ww,nH
-        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wd*Wh*Ww, Wd*Wh*Ww
-        attn = attn + relative_position_bias.unsqueeze(0) # B_, nH, N, N
-
-        if mask is not None:
-            nW = mask.shape[0]
-            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
-            attn = attn.view(-1, self.num_heads, N, N)
-            attn = self.softmax(attn)
-        else:
-            attn = self.softmax(attn)
-
-        attn = self.attn_drop(attn)
-        # print('attn: ', attn.shape, ', v: ', v.shape, ', x: ', x.shape)
-        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
-        x = self.proj(x)
-        x = self.proj_drop(x)
-        return x
-
-
-class SwinTransformerBlock3D(nn.Module):
-    """ Swin Transformer Block.
-    Args:
-        dim (int): Number of input channels.
-        num_heads (int): Number of attention heads.
-        window_size (tuple[int]): Window size.
-        shift_size (tuple[int]): Shift size for SW-MSA.
-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
-        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
-        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
-        drop (float, optional): Dropout rate. Default: 0.0
-        attn_drop (float, optional): Attention dropout rate. Default: 0.0
-        drop_path (float, optional): Stochastic depth rate. Default: 0.0
-        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
-        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
-    """
+    def forward(self, c_s, h_s, t):
+        # print(c_s.size(), h_s.size(), t.size())
+        phi = self._compute_phi(t)
 
-    def __init__(self, dim, num_heads, window_size=(2,7,7), shift_size=(0,0,0),
-                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
-                 act_layer=nn.GELU, norm_layer=nn.LayerNorm, use_checkpoint=False):
-        super().__init__()
-        self.dim = dim
-        self.num_heads = num_heads
-        self.window_size = window_size
-        self.shift_size = shift_size
-        self.mlp_ratio = mlp_ratio
-        self.use_checkpoint=use_checkpoint
-
-        assert 0 <= self.shift_size[0] < self.window_size[0], "shift_size must in 0-window_size"
-        assert 0 <= self.shift_size[1] < self.window_size[1], "shift_size must in 0-window_size"
-        assert 0 <= self.shift_size[2] < self.window_size[2], "shift_size must in 0-window_size"
-
-        self.norm1 = norm_layer(dim)
-        self.attn = WindowAttention3D(
-            dim, window_size=self.window_size, num_heads=num_heads,
-            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
-
-        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
-        self.norm2 = norm_layer(dim)
-        mlp_hidden_dim = int(dim * mlp_ratio)
-        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
-
-    def forward_part1(self, x, mask_matrix):
-        B, D, H, W, C = x.shape
-        window_size, shift_size = get_window_size((D, H, W), self.window_size, self.shift_size)
-        # print('window_size: ', window_size, ', shift_size: ', shift_size)
-        x = self.norm1(x)
-        # pad feature maps to multiples of window size
-        pad_l = pad_t = pad_d0 = 0
-        pad_d1 = (window_size[0] - D % window_size[0]) % window_size[0]
-        pad_b = (window_size[1] - H % window_size[1]) % window_size[1]
-        pad_r = (window_size[2] - W % window_size[2]) % window_size[2]
-        x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b, pad_d0, pad_d1))
-        _, Dp, Hp, Wp, _ = x.shape
-        # cyclic shift
-        if any(i > 0 for i in shift_size):
-            shifted_x = torch.roll(x, shifts=(-shift_size[0], -shift_size[1], -shift_size[2]), dims=(1, 2, 3))
-            attn_mask = mask_matrix
-        else:
-            shifted_x = x
-            attn_mask = None
-        # partition windows
-        x_windows = window_partition(shifted_x, window_size)  # B*nW, Wd*Wh*Ww, C
-        # print('shifted_x: ', shifted_x.shape, 'x_windows: ', x_windows.shape)
-        # W-MSA/SW-MSA
-        attn_windows = self.attn(x_windows, mask=attn_mask)  # B*nW, Wd*Wh*Ww, C
-        # merge windows
-        attn_windows = attn_windows.view(-1, *(window_size+(C,)))
-        # print('attn_windows: ', attn_windows.shape)
-        shifted_x = window_reverse(attn_windows, window_size, B, Dp, Hp, Wp)  # B D' H' W' C
-        # reverse cyclic shift
-        if any(i > 0 for i in shift_size):
-            x = torch.roll(shifted_x, shifts=(shift_size[0], shift_size[1], shift_size[2]), dims=(1, 2, 3))
-        else:
-            x = shifted_x
+        # Phase-related augmentations
+        k_up = 2 * phi / self.ratio_on
+        k_down = 2 - k_up
+        k_closed = self.leak * phi
 
-        if pad_d1 >0 or pad_r > 0 or pad_b > 0:
-            x = x[:, :D, :H, :W, :].contiguous()
-        return x
+        k = torch.where(phi < self.ratio_on, k_down, k_closed)
+        k = torch.where(phi < 0.5 * self.ratio_on, k_up, k)
+        k = k.view(c_s.shape[0], -1)
+        c_s_new = k * c_s + (1 - k) * self.c0
+        h_s_new = k * h_s + (1 - k) * self.h0
 
-    def forward_part2(self, x):
-        return self.drop_path(self.mlp(self.norm2(x)))
+        return h_s_new, c_s_new
 
-    def forward(self, x, mask_matrix):
-        """ Forward function.
-        Args:
-            x: Input feature, tensor size (B, D, H, W, C).
-            mask_matrix: Attention mask for cyclic shift.
-        """
 
-        shortcut = x
-        if self.use_checkpoint:
-            x = checkpoint.checkpoint(self.forward_part1, x, mask_matrix)
-        else:
-            x = self.forward_part1(x, mask_matrix)
-        x = shortcut + self.drop_path(x)
+class ConvLSTM(nn.Module):
+    """Adapted from: https://github.com/Atcold/pytorch-CortexNet/blob/master/model/ConvLSTMCell.py """
 
-        if self.use_checkpoint:
-            x = x + checkpoint.checkpoint(self.forward_part2, x)
-        else:
-            x = x + self.forward_part2(x)
+    def __init__(self, input_size, hidden_size, kernel_size):
+        super(ConvLSTM, self).__init__()
 
-        return x
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        pad = kernel_size // 2
 
+        # cache a tensor filled with zeros to avoid reallocating memory at each inference step if --no-recurrent is enabled
+        self.zero_tensors = {}
 
-class PatchMerging(nn.Module):
-    """ Patch Merging Layer
-    Args:
-        dim (int): Number of input channels.
-        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
-    """
-    def __init__(self, dim, norm_layer=nn.LayerNorm):
-        super().__init__()
-        self.dim = dim
-        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
-        self.norm = norm_layer(4 * dim)
+        self.Gates = nn.Conv2d(input_size + hidden_size, 4 * hidden_size, kernel_size, padding=pad)
 
-    def forward(self, x):
-        """ Forward function.
-        Args:
-            x: Input feature, tensor size (B, D, H, W, C).
-        """
-        B, D, H, W, C = x.shape
-
-        # padding
-        pad_input = (H % 2 == 1) or (W % 2 == 1)
-        if pad_input:
-            x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2))
-
-        x0 = x[:, :, 0::2, 0::2, :]  # B D H/2 W/2 C
-        x1 = x[:, :, 1::2, 0::2, :]  # B D H/2 W/2 C
-        x2 = x[:, :, 0::2, 1::2, :]  # B D H/2 W/2 C
-        x3 = x[:, :, 1::2, 1::2, :]  # B D H/2 W/2 C
-        x = torch.cat([x0, x1, x2, x3], -1)  # B D H/2 W/2 4*C
-
-        x = self.norm(x)
-        x = self.reduction(x)
-
-        return x
-
-
-# cache each stage results
-@lru_cache()
-def compute_mask(D, H, W, window_size, shift_size, device):
-    img_mask = torch.zeros((1, D, H, W, 1), device=device)  # 1 Dp Hp Wp 1
-    cnt = 0
-    for d in slice(-window_size[0]), slice(-window_size[0], -shift_size[0]), slice(-shift_size[0],None):
-        for h in slice(-window_size[1]), slice(-window_size[1], -shift_size[1]), slice(-shift_size[1],None):
-            for w in slice(-window_size[2]), slice(-window_size[2], -shift_size[2]), slice(-shift_size[2],None):
-                img_mask[:, d, h, w, :] = cnt
-                cnt += 1
-    mask_windows = window_partition(img_mask, window_size)  # nW, ws[0]*ws[1]*ws[2], 1
-    mask_windows = mask_windows.squeeze(-1)  # nW, ws[0]*ws[1]*ws[2]
-    attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
-    attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
-    return attn_mask
-
-
-class BasicLayer(nn.Module):
-    """ A basic Swin Transformer layer for one stage.
-    Args:
-        dim (int): Number of feature channels
-        depth (int): Depths of this stage.
-        num_heads (int): Number of attention head.
-        window_size (tuple[int]): Local window size. Default: (1,7,7).
-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
-        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
-        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
-        drop (float, optional): Dropout rate. Default: 0.0
-        attn_drop (float, optional): Attention dropout rate. Default: 0.0
-        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
-        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
-        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
-    """
+    def forward(self, input_, prev_state=None):
+        # get batch and spatial sizes
+        batch_size = input_.data.size()[0]
+        spatial_size = input_.data.size()[2:]
 
-    def __init__(self,
-                 dim,
-                 depth,
-                 num_heads,
-                 window_size=(1,7,7),
-                 mlp_ratio=4.,
-                 qkv_bias=False,
-                 qk_scale=None,
-                 drop=0.,
-                 attn_drop=0.,
-                 drop_path=0.,
-                 norm_layer=nn.LayerNorm,
-                 downsample=None,
-                 use_checkpoint=False):
-        super().__init__()
-        self.window_size = window_size
-        self.shift_size = tuple(i // 2 for i in window_size)
-        self.depth = depth
-        self.use_checkpoint = use_checkpoint
-
-        # build blocks
-        self.blocks = nn.ModuleList([
-            SwinTransformerBlock3D(
-                dim=dim,
-                num_heads=num_heads,
-                window_size=window_size,
-                shift_size=(0,0,0) if (i % 2 == 0) else self.shift_size,
-                mlp_ratio=mlp_ratio,
-                qkv_bias=qkv_bias,
-                qk_scale=qk_scale,
-                drop=drop,
-                attn_drop=attn_drop,
-                drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
-                norm_layer=norm_layer,
-                use_checkpoint=use_checkpoint,
-            )
-            for i in range(depth)])
-        
-        self.downsample = downsample
-        if self.downsample is not None:
-            self.downsample = downsample(dim=dim, norm_layer=norm_layer)
+        # generate empty prev_state, if None is provided
+        if prev_state is None:
 
-    def forward(self, x):
-        """ Forward function.
-        Args:
-            x: Input feature, tensor size (B, C, D, H, W).
-        """
-        # calculate attention mask for SW-MSA
-        B, C, D, H, W = x.shape
-        window_size, shift_size = get_window_size((D,H,W), self.window_size, self.shift_size)
-        x = rearrange(x, 'b c d h w -> b d h w c')
-        Dp = int(np.ceil(D / window_size[0])) * window_size[0]
-        Hp = int(np.ceil(H / window_size[1])) * window_size[1]
-        Wp = int(np.ceil(W / window_size[2])) * window_size[2]
-        attn_mask = compute_mask(Dp, Hp, Wp, window_size, shift_size, x.device)
-        for blk in self.blocks:
-            x = blk(x, attn_mask)
-        # print(x.shape)
-        x = x.view(B, D, H, W, -1)
-
-        if self.downsample is not None:
-            x_out = self.downsample(x)
-        else:
-            x_out = x
-        x_out = rearrange(x_out, 'b d h w c -> b c d h w')
-        return x_out, x
-
-
-class PatchEmbed3D(nn.Module):
-    """ Video to Patch Embedding.
-    Args:
-        patch_size (int): Patch token size. Default: (2,4,4).
-        in_chans (int): Number of input video channels. Default: 3.
-        embed_dim (int): Number of linear projection output channels. Default: 96.
-        norm_layer (nn.Module, optional): Normalization layer. Default: None
-    """
-    def __init__(self, patch_size=(2,4,4), in_chans=3, embed_dim=96, norm_layer=None):
-        super().__init__()
-        self.patch_size = patch_size
+            # create the zero tensor if it has not been created already
+            state_size = tuple([batch_size, self.hidden_size] + list(spatial_size))
+            if state_size not in self.zero_tensors:
+                # allocate a tensor with size `spatial_size`, filled with zero (if it has not been allocated already)
+                self.zero_tensors[state_size] = (
+                    torch.zeros(state_size, dtype=input_.dtype).to(input_.device),
+                    torch.zeros(state_size, dtype=input_.dtype).to(input_.device)
+                )
+
+            prev_state = self.zero_tensors[tuple(state_size)]
+
+        prev_hidden, prev_cell = prev_state
 
-        self.in_chans = in_chans
-        self.embed_dim = embed_dim
+        # data size is [batch, channel, height, width]
+        stacked_inputs = torch.cat((input_, prev_hidden), 1)
+        gates = self.Gates(stacked_inputs)
 
-        # self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
-        self.proj = nn.Conv3d(1, embed_dim, kernel_size=patch_size, stride=patch_size)
-        if norm_layer is not None:
-            self.norm = norm_layer(embed_dim)
+        # chunk across channel dimension
+        in_gate, remember_gate, out_gate, cell_gate = gates.chunk(4, 1)
+
+        # apply sigmoid non linearity
+        in_gate = torch.sigmoid(in_gate)
+        remember_gate = torch.sigmoid(remember_gate)
+        out_gate = torch.sigmoid(out_gate)
+
+        # apply tanh non linearity
+        cell_gate = torch.tanh(cell_gate)
+
+        # compute current cell and hidden state
+        cell = (remember_gate * prev_cell) + (in_gate * cell_gate)
+        hidden = out_gate * torch.tanh(cell)
+
+        return hidden, cell
+
+
+class PhasedConvLSTMCell(nn.Module):
+    def __init__(
+        self,
+        input_channels,
+        hidden_channels,
+        kernel_size=3
+    ):
+        super().__init__()
+        self.hidden_channels = hidden_channels
+
+        self.lstm = ConvLSTM(
+            input_size=input_channels,
+            hidden_size=hidden_channels,
+            kernel_size=kernel_size
+        )
+
+        # as soon as spatial dimension is known, phased lstm cell is instantiated
+        self.phased_cell = None
+        self.hidden_size = None
+
+    def forward(self, input, times, prev_state=None):
+        # input: B x C x H x W
+        # times: B
+        # returns: output: B x C_out x H x W,   prev_state: (B x C_out x H x W, B x C_out x H x W)
+
+        B, C, H, W = input.shape
+
+        if self.hidden_size is None:
+            self.hidden_size = self.hidden_channels * W * H
+            self.phased_cell = PhasedLSTMCell(hidden_size=self.hidden_size)
+            self.phased_cell = self.phased_cell.to(input.device)
+            self.phased_cell.requires_grad = True
+
+        if prev_state is None:
+            h0 = input.new_zeros((B, self.hidden_channels, H, W))
+            c0 = input.new_zeros((B, self.hidden_channels, H, W))
         else:
-            self.norm = None
+            c0, h0 = prev_state
 
-    def forward(self, x):
-        """Forward function."""
-        # padding
-        x = x.unsqueeze(1)
-        _, _, D, H, W = x.size()
-        if W % self.patch_size[2] != 0:
-            x = F.pad(x, (0, self.patch_size[2] - W % self.patch_size[2]))
-        if H % self.patch_size[1] != 0:
-            x = F.pad(x, (0, 0, 0, self.patch_size[1] - H % self.patch_size[1]))
-        if D % self.patch_size[0] != 0:
-            x = F.pad(x, (0, 0, 0, 0, 0, self.patch_size[0] - D % self.patch_size[0]))
-
-        x = self.proj(x)  # B C D Wh Ww
-        if self.norm is not None:
-            D, Wh, Ww = x.size(2), x.size(3), x.size(4)
-            x = x.flatten(2).transpose(1, 2)
-            x = self.norm(x)
-            x = x.transpose(1, 2).view(-1, self.embed_dim, D, Wh, Ww)
-
-        return x
-
-
-class SwinTransformer3D(nn.Module):
-    """ Swin Transformer backbone.
-    Args:
-        patch_size (int | tuple(int)): Patch size. Default: (4,4,4).
-        in_chans (int): Number of input image channels. Default: 3.
-        embed_dim (int): Number of linear projection output channels. Default: 96.
-        depths (tuple[int]): Depths of each Swin Transformer stage.
-        num_heads (tuple[int]): Number of attention head of each stage.
-        window_size (int): Window size. Default: 7.
-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
-        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: Truee
-        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set.
-        drop_rate (float): Dropout rate.
-        attn_drop_rate (float): Attention dropout rate. Default: 0.
-        drop_path_rate (float): Stochastic depth rate. Default: 0.2.
-        norm_layer: Normalization layer. Default: nn.LayerNorm.
-        patch_norm (bool): If True, add normalization after patch embedding. Default: False.
-        frozen_stages (int): Stages to be frozen (stop grad and set eval mode).
-            -1 means not freezing any parameters.
+        self.phased_cell.set_state(c0.view(B, -1), h0.view(B, -1))
+
+        c_t, h_t = self.lstm(input, (c0, h0))
+
+        # reshape activation maps such that phased lstm can use them
+        (c_s, h_s) = self.phased_cell(c_t.view(B, -1), h_t.view(B, -1), times)
+
+        # reshape to feed to conv lstm
+        c_s = c_s.view(B, -1, H, W)
+        h_s = h_s.view(B, -1, H, W)
+
+        return h_t, (c_s, h_s)
+
+
+class ConvGRU(nn.Module):
+    """
+    Generate a convolutional GRU cell
+    Adapted from: https://github.com/jacobkimmel/pytorch_convgru/blob/master/convgru.py
     """
 
-    def __init__(self,
-                 pretrained=None,
-                 pretrained2d=True,
-                 patch_size=(4,4,4),
-                 in_chans=3,
-                 embed_dim=96,
-                 depths=[2, 2, 6, 2],
-                 num_heads=[3, 6, 12, 24],
-                 window_size=(2,7,7),
-                 mlp_ratio=4.,
-                 qkv_bias=True,
-                 qk_scale=None,
-                 drop_rate=0.,
-                 attn_drop_rate=0.,
-                 drop_path_rate=0.2,
-                 norm_layer=nn.LayerNorm,
-                 patch_norm=False,
-                 out_indices=(0,1,2,3),
-                 frozen_stages=-1,
-                 use_checkpoint=False,
-                 new_version=0):
+    def __init__(self, input_size, hidden_size, kernel_size):
         super().__init__()
-
-        self.pretrained = pretrained
-        self.pretrained2d = pretrained2d
-        self.num_layers = len(depths)
-        self.embed_dim = embed_dim
-        self.patch_norm = patch_norm
-        self.frozen_stages = frozen_stages
-        self.window_size = window_size
-        self.patch_size = patch_size
-        self.out_indices = out_indices
-
-        # split image into non-overlapping patches
-        if new_version==3:
-            print("---- new version 3 ----")
-            self.patch_embed = PatchEmbedConv(
-                patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,
-                norm_layer=norm_layer if self.patch_norm else None)
-        elif new_version==4:
-            print("---- new version 4 ----")
-            self.patch_embed = PatchEmbedLocalGlobal(
-                patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,
-                norm_layer=norm_layer if self.patch_norm else None)
+        padding = kernel_size // 2
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.reset_gate = nn.Conv2d(input_size + hidden_size, hidden_size, kernel_size, padding=padding)
+        self.update_gate = nn.Conv2d(input_size + hidden_size, hidden_size, kernel_size, padding=padding)
+        self.out_gate = nn.Conv2d(input_size + hidden_size, hidden_size, kernel_size, padding=padding)
+
+        init.orthogonal_(self.reset_gate.weight)
+        init.orthogonal_(self.update_gate.weight)
+        init.orthogonal_(self.out_gate.weight)
+        init.constant_(self.reset_gate.bias, 0.)
+        init.constant_(self.update_gate.bias, 0.)
+        init.constant_(self.out_gate.bias, 0.)
+
+    def forward(self, input_, prev_state):
+
+        # get batch and spatial sizes
+        batch_size = input_.data.size()[0]
+        spatial_size = input_.data.size()[2:]
+
+        # generate empty prev_state, if None is provided
+        if prev_state is None:
+            state_size = [batch_size, self.hidden_size] + list(spatial_size)
+            prev_state = torch.zeros(state_size, dtype=input_.dtype).to(input_.device)
+
+        # data size is [batch, channel, height, width]
+        stacked_inputs = torch.cat([input_, prev_state], dim=1)
+        update = torch.sigmoid(self.update_gate(stacked_inputs))
+        reset = torch.sigmoid(self.reset_gate(stacked_inputs))
+        out_inputs = torch.tanh(self.out_gate(torch.cat([input_, prev_state * reset], dim=1)))
+        new_state = prev_state * (1 - update) + out_inputs * update
+
+        return new_state
+
+
+class RecurrentResidualLayer(nn.Module):
+    def __init__(self, in_channels, out_channels,
+                 recurrent_block_type='convlstm', norm=None, BN_momentum=0.1):
+        super(RecurrentResidualLayer, self).__init__()
+
+        assert(recurrent_block_type in ['convlstm', 'convgru'])
+        self.recurrent_block_type = recurrent_block_type
+        if self.recurrent_block_type == 'convlstm':
+            RecurrentBlock = ConvLSTM
         else:
-            print("---- old version ----")
-            self.patch_embed = PatchEmbed3D(
-                patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,
-                norm_layer=norm_layer if self.patch_norm else None)
-
-        self.pos_drop = nn.Dropout(p=drop_rate)
-
-        # stochastic depth
-        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule
-
-        # build layers
-        self.layers = nn.ModuleList()
-        for i_layer in range(self.num_layers):
-            layer = BasicLayer(
-                dim=int(embed_dim * 2**i_layer),
-                depth=depths[i_layer],
-                num_heads=num_heads[i_layer],
-                window_size=window_size,
-                mlp_ratio=mlp_ratio,
-                qkv_bias=qkv_bias,
-                qk_scale=qk_scale,
-                drop=drop_rate,
-                attn_drop=attn_drop_rate,
-                drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
-                norm_layer=norm_layer,
-                downsample=PatchMerging if i_layer<self.num_layers-1 else None,
-                use_checkpoint=use_checkpoint)
-            self.layers.append(layer)
-
-        # self.num_features = int(embed_dim * 2**(self.num_layers-1))
-        num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)]
-        self.num_features = num_features
-
-        # add a norm layer for each output
-        # self.norm = norm_layer(self.num_features)
-
-        # add a norm layer for each output
-        for i_layer in self.out_indices:
-            layer = norm_layer(self.num_features[i_layer])
-            layer_name = f'norm{i_layer}'
-            self.add_module(layer_name, layer)
-
-
-    def inflate_weights(self, logger):
-        """Inflate the swin2d parameters to swin3d.
-        The differences between swin3d and swin2d mainly lie in an extra
-        axis. To utilize the pretrained parameters in 2d model,
-        the weight of swin2d models should be inflated to fit in the shapes of
-        the 3d counterpart.
-        Args:
-            logger (logging.Logger): The logger used to print
-                debugging infomation.
-        """
-        checkpoint = torch.load(self.pretrained, map_location='cpu')
-        state_dict = checkpoint['model']
-
-        # delete relative_position_index since we always re-init it
-        relative_position_index_keys = [k for k in state_dict.keys() if "relative_position_index" in k]
-        for k in relative_position_index_keys:
-            del state_dict[k]
-
-        # delete attn_mask since we always re-init it
-        attn_mask_keys = [k for k in state_dict.keys() if "attn_mask" in k]
-        for k in attn_mask_keys:
-            del state_dict[k]
-
-        state_dict['patch_embed.proj.weight'] = state_dict['patch_embed.proj.weight'].unsqueeze(2).repeat(1,1,self.patch_size[0],1,1) / self.patch_size[0]
-
-        # bicubic interpolate relative_position_bias_table if not match
-        relative_position_bias_table_keys = [k for k in state_dict.keys() if "relative_position_bias_table" in k]
-        for k in relative_position_bias_table_keys:
-            relative_position_bias_table_pretrained = state_dict[k]
-            relative_position_bias_table_current = self.state_dict()[k]
-            L1, nH1 = relative_position_bias_table_pretrained.size()
-            L2, nH2 = relative_position_bias_table_current.size()
-            L2 = (2*self.window_size[1]-1) * (2*self.window_size[2]-1)
-            wd = self.window_size[0]
-            if nH1 != nH2:
-                logger.warning(f"Error in loading {k}, passing")
-            else:
-                if L1 != L2:
-                    S1 = int(L1 ** 0.5)
-                    relative_position_bias_table_pretrained_resized = torch.nn.functional.interpolate(
-                        relative_position_bias_table_pretrained.permute(1, 0).view(1, nH1, S1, S1), size=(2*self.window_size[1]-1, 2*self.window_size[2]-1),
-                        mode='bicubic')
-                    relative_position_bias_table_pretrained = relative_position_bias_table_pretrained_resized.view(nH2, L2).permute(1, 0)
-            state_dict[k] = relative_position_bias_table_pretrained.repeat(2*wd-1,1)
-
-        msg = self.load_state_dict(state_dict, strict=False)
-        logger.info(msg)
-        logger.info(f"=> loaded successfully '{self.pretrained}'")
-        del checkpoint
-        torch.cuda.empty_cache()
-
-    def init_weights(self, pretrained=None):
-        """Initialize the weights in backbone.
-        Args:
-            pretrained (str, optional): Path to pre-trained weights.
-                Defaults to None.
-        """
-        def _init_weights(m):
-            if isinstance(m, nn.Linear):
-                trunc_normal_(m.weight, std=.02)
-                if isinstance(m, nn.Linear) and m.bias is not None:
-                    nn.init.constant_(m.bias, 0)
-            elif isinstance(m, nn.LayerNorm):
-                nn.init.constant_(m.bias, 0)
-                nn.init.constant_(m.weight, 1.0)
-
-        if pretrained:
-            self.pretrained = pretrained
-        if isinstance(self.pretrained, str):
-            self.apply(_init_weights)
-            logger = get_root_logger()
-            logger.info(f'load model from: {self.pretrained}')
-
-            if self.pretrained2d:
-                # Inflate 2D model into 3D model.
-                self.inflate_weights(logger)
-            else:
-                # Directly load 3D model.
-                load_checkpoint(self, self.pretrained, strict=False, logger=logger)
-        elif self.pretrained is None:
-            self.apply(_init_weights)
-        else:
-            raise TypeError('pretrained must be a str or None')
-
-    def forward(self, x):
-        """Forward function."""
-        x = self.patch_embed(x)
-        # print(x.shape)
-        x = self.pos_drop(x)
-        
-        outs = []
-        for i, layer in enumerate(self.layers):
-            x, out_x = layer(x.contiguous())
-            # print('---- ', out_x.shape)
-            if i in self.out_indices:
-                norm_layer = getattr(self, f'norm{i}')
-                out_x = norm_layer(out_x)
-                _, Ti, Hi, Wi, Ci = out_x.shape 
-                out = rearrange(out_x, 'n d h w c -> n c d h w')
-                outs.append(out)
-
-        return tuple(outs)
+            RecurrentBlock = ConvGRU
+        self.conv = ResidualBlock(in_channels=in_channels,
+                                  out_channels=out_channels,
+                                  norm=norm,
+                                  BN_momentum=BN_momentum)
+        self.recurrent_block = RecurrentBlock(input_size=out_channels,
+                                              hidden_size=out_channels,
+                                              kernel_size=3)
+
+    def forward(self, x, prev_state):
+        x = self.conv(x)
+        state = self.recurrent_block(x, prev_state)
+        x = state[0] if self.recurrent_block_type == 'convlstm' else state
+        return x, state