| 161 | |
| 162 | |
| 163 | class SplitSABlock(nn.Module): |
| 164 | def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., |
| 165 | drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm): |
| 166 | super().__init__() |
| 167 | self.pos_embed = conv_3x3x3(dim, dim, groups=dim) |
| 168 | self.t_norm = norm_layer(dim) |
| 169 | self.t_attn = Attention( |
| 170 | dim, |
| 171 | num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, |
| 172 | attn_drop=attn_drop, proj_drop=drop) |
| 173 | self.norm1 = norm_layer(dim) |
| 174 | self.attn = Attention( |
| 175 | dim, |
| 176 | num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, |
| 177 | attn_drop=attn_drop, proj_drop=drop) |
| 178 | # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here |
| 179 | self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() |
| 180 | self.norm2 = norm_layer(dim) |
| 181 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 182 | self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) |
| 183 | |
| 184 | def forward(self, x): |
| 185 | x = x + self.pos_embed(x) |
| 186 | B, C, T, H, W = x.shape |
| 187 | attn = x.view(B, C, T, H * W).permute(0, 3, 2, 1).contiguous() |
| 188 | attn = attn.view(B * H * W, T, C) |
| 189 | attn = attn + self.drop_path(self.t_attn(self.t_norm(attn))) |
| 190 | attn = attn.view(B, H * W, T, C).permute(0, 2, 1, 3).contiguous() |
| 191 | attn = attn.view(B * T, H * W, C) |
| 192 | residual = x.view(B, C, T, H * W).permute(0, 2, 3, 1).contiguous() |
| 193 | residual = residual.view(B * T, H * W, C) |
| 194 | attn = residual + self.drop_path(self.attn(self.norm1(attn))) |
| 195 | attn = attn.view(B, T * H * W, C) |
| 196 | out = attn + self.drop_path(self.mlp(self.norm2(attn))) |
| 197 | out = out.transpose(1, 2).reshape(B, C, T, H, W) |
| 198 | return out |
| 199 | |
| 200 | |
| 201 | class SpeicalPatchEmbed(nn.Module): |