| 135 | |
| 136 | |
| 137 | class SABlock(nn.Module): |
| 138 | def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., |
| 139 | drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm): |
| 140 | super().__init__() |
| 141 | self.pos_embed = conv_3x3x3(dim, dim, groups=dim) |
| 142 | self.norm1 = norm_layer(dim) |
| 143 | self.attn = Attention( |
| 144 | dim, |
| 145 | num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, |
| 146 | attn_drop=attn_drop, proj_drop=drop) |
| 147 | # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here |
| 148 | self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() |
| 149 | self.norm2 = norm_layer(dim) |
| 150 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 151 | self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) |
| 152 | |
| 153 | def forward(self, x): |
| 154 | x = x + self.pos_embed(x) |
| 155 | B, C, T, H, W = x.shape |
| 156 | x = x.flatten(2).transpose(1, 2) |
| 157 | x = x + self.drop_path(self.attn(self.norm1(x))) |
| 158 | x = x + self.drop_path(self.mlp(self.norm2(x))) |
| 159 | x = x.transpose(1, 2).reshape(B, C, T, H, W) |
| 160 | return x |
| 161 | |
| 162 | |
| 163 | class SplitSABlock(nn.Module): |