Parallel block class.
| 229 | |
| 230 | |
| 231 | class ParallelBlock(nn.Module): |
| 232 | """ Parallel block class. """ |
| 233 | def __init__(self, dims, num_heads, mlp_ratios=[], qkv_bias=False, drop=0., attn_drop=0., |
| 234 | drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, shared_crpes=None): |
| 235 | super().__init__() |
| 236 | |
| 237 | # Conv-Attention. |
| 238 | self.norm12 = norm_layer(dims[1]) |
| 239 | self.norm13 = norm_layer(dims[2]) |
| 240 | self.norm14 = norm_layer(dims[3]) |
| 241 | self.factoratt_crpe2 = FactorAtt_ConvRelPosEnc( |
| 242 | dims[1], num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop, |
| 243 | shared_crpe=shared_crpes[1] |
| 244 | ) |
| 245 | self.factoratt_crpe3 = FactorAtt_ConvRelPosEnc( |
| 246 | dims[2], num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop, |
| 247 | shared_crpe=shared_crpes[2] |
| 248 | ) |
| 249 | self.factoratt_crpe4 = FactorAtt_ConvRelPosEnc( |
| 250 | dims[3], num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop, |
| 251 | shared_crpe=shared_crpes[3] |
| 252 | ) |
| 253 | self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() |
| 254 | |
| 255 | # MLP. |
| 256 | self.norm22 = norm_layer(dims[1]) |
| 257 | self.norm23 = norm_layer(dims[2]) |
| 258 | self.norm24 = norm_layer(dims[3]) |
| 259 | # In parallel block, we assume dimensions are the same and share the linear transformation. |
| 260 | assert dims[1] == dims[2] == dims[3] |
| 261 | assert mlp_ratios[1] == mlp_ratios[2] == mlp_ratios[3] |
| 262 | mlp_hidden_dim = int(dims[1] * mlp_ratios[1]) |
| 263 | self.mlp2 = self.mlp3 = self.mlp4 = Mlp( |
| 264 | in_features=dims[1], hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) |
| 265 | |
| 266 | def upsample(self, x, factor: float, size: Tuple[int, int]): |
| 267 | """ Feature map up-sampling. """ |
| 268 | return self.interpolate(x, scale_factor=factor, size=size) |
| 269 | |
| 270 | def downsample(self, x, factor: float, size: Tuple[int, int]): |
| 271 | """ Feature map down-sampling. """ |
| 272 | return self.interpolate(x, scale_factor=1.0/factor, size=size) |
| 273 | |
| 274 | def interpolate(self, x, scale_factor: float, size: Tuple[int, int]): |
| 275 | """ Feature map interpolation. """ |
| 276 | B, N, C = x.shape |
| 277 | H, W = size |
| 278 | assert N == 1 + H * W |
| 279 | |
| 280 | cls_token = x[:, :1, :] |
| 281 | img_tokens = x[:, 1:, :] |
| 282 | |
| 283 | img_tokens = img_tokens.transpose(1, 2).reshape(B, C, H, W) |
| 284 | img_tokens = F.interpolate( |
| 285 | img_tokens, scale_factor=scale_factor, recompute_scale_factor=False, mode='bilinear', align_corners=False) |
| 286 | img_tokens = img_tokens.reshape(B, C, -1).transpose(1, 2) |
| 287 | |
| 288 | out = torch.cat((cls_token, img_tokens), dim=1) |