Args: depth (int): Number of blocks. num_heads (int): Number of attention heads. drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm use_checkpoint (
| 17 | |
| 18 | |
| 19 | class ResidualBlock(nn.Module): |
| 20 | """ |
| 21 | |
| 22 | Args: |
| 23 | depth (int): Number of blocks. |
| 24 | num_heads (int): Number of attention heads. |
| 25 | drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 |
| 26 | norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm |
| 27 | use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. |
| 28 | resi_connection: The convolutional block before residual connection. |
| 29 | """ |
| 30 | |
| 31 | def __init__(self, embed_dim, embed_dim_next, depth, num_heads, heads_range, init_value, |
| 32 | ffn_dim=96, layerscale=False, layer_init_values=1e-5, |
| 33 | drop_path=0., norm_layer=nn.LayerNorm, use_checkpoint=False, |
| 34 | resi_connection='1conv', use_pos_map=False): |
| 35 | super(ResidualBlock, self).__init__() |
| 36 | |
| 37 | self.embed_dim = embed_dim |
| 38 | |
| 39 | self.residual_group = BasicLayer(embed_dim=embed_dim, |
| 40 | depth=depth, |
| 41 | num_heads=num_heads, |
| 42 | heads_range=heads_range, |
| 43 | init_value=init_value, |
| 44 | drop_path=drop_path, |
| 45 | ffn_dim=ffn_dim, |
| 46 | norm_layer=norm_layer, |
| 47 | use_checkpoint=use_checkpoint, |
| 48 | layerscale=layerscale, |
| 49 | layer_init_values=layer_init_values, |
| 50 | ) |
| 51 | |
| 52 | embed_dim_conv_in = embed_dim + 2 if use_pos_map else embed_dim |
| 53 | if resi_connection == '1conv': |
| 54 | self.conv = nn.Conv2d(embed_dim_conv_in, embed_dim, 3, 1, 1) |
| 55 | elif resi_connection == '3conv': |
| 56 | # to save parameters and memory |
| 57 | self.conv = nn.Sequential(nn.Conv2d(embed_dim_conv_in, embed_dim // 4, 3, 1, 1), |
| 58 | nn.LeakyReLU(negative_slope=0.2, inplace=True), |
| 59 | nn.Conv2d(embed_dim // 4, embed_dim // 4, 1, 1, 0), |
| 60 | nn.LeakyReLU(negative_slope=0.2, inplace=True), |
| 61 | nn.Conv2d(embed_dim // 4, embed_dim, 3, 1, 1)) |
| 62 | |
| 63 | self.patch_embed = PatchEmbedIR(embed_dim=embed_dim) |
| 64 | |
| 65 | self.patch_unembed = PatchUnEmbedIR(embed_dim=embed_dim) |
| 66 | |
| 67 | self.final_op = ChannelEmbeddingCompression(embed_dim, embed_dim_next) if ( |
| 68 | embed_dim != embed_dim_next) else IdentityMod() |
| 69 | |
| 70 | self.use_pos_map = use_pos_map |
| 71 | |
| 72 | def forward(self, x, x_size, pos_map=None, att_range_factor=None, **kwargs): |
| 73 | |
| 74 | rb_stl_out = self.residual_group(x, att_range_factor) |
| 75 | |
| 76 | rb_unembed = self.patch_unembed(rb_stl_out, x_size) |