r""" Focal Modulation Network Block. Args: dim (int): Number of input channels. input_resolution (tuple[int]): Input resulotion. mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. drop (float, optional): Dropout rate. Default: 0.0 drop_path (flo
| 125 | |
| 126 | |
| 127 | class FocalNetBlock(nn.Module): |
| 128 | r""" Focal Modulation Network Block. |
| 129 | |
| 130 | Args: |
| 131 | dim (int): Number of input channels. |
| 132 | input_resolution (tuple[int]): Input resulotion. |
| 133 | mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. |
| 134 | drop (float, optional): Dropout rate. Default: 0.0 |
| 135 | drop_path (float, optional): Stochastic depth rate. Default: 0.0 |
| 136 | act_layer (nn.Module, optional): Activation layer. Default: nn.GELU |
| 137 | norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm |
| 138 | focal_level (int): Number of focal levels. |
| 139 | focal_window (int): Focal window size at first focal level |
| 140 | use_layerscale (bool): Whether use layerscale |
| 141 | layerscale_value (float): Initial layerscale value |
| 142 | use_postln (bool): Whether use layernorm after modulation |
| 143 | """ |
| 144 | |
| 145 | def __init__(self, dim, input_resolution, mlp_ratio=4., drop=0., drop_path=0., |
| 146 | act_layer=nn.GELU, norm_layer=nn.LayerNorm, |
| 147 | focal_level=1, focal_window=3, |
| 148 | use_layerscale=False, layerscale_value=1e-4, |
| 149 | use_postln=False, use_postln_in_modulation=False, |
| 150 | normalize_modulator=False): |
| 151 | super().__init__() |
| 152 | self.dim = dim |
| 153 | self.input_resolution = input_resolution |
| 154 | self.mlp_ratio = mlp_ratio |
| 155 | |
| 156 | self.focal_window = focal_window |
| 157 | self.focal_level = focal_level |
| 158 | self.use_postln = use_postln |
| 159 | |
| 160 | self.norm1 = norm_layer(dim) |
| 161 | self.modulation = FocalModulation( |
| 162 | dim, proj_drop=drop, focal_window=focal_window, focal_level=self.focal_level, |
| 163 | use_postln_in_modulation=use_postln_in_modulation, normalize_modulator=normalize_modulator |
| 164 | ) |
| 165 | |
| 166 | self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() |
| 167 | self.norm2 = norm_layer(dim) |
| 168 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 169 | self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) |
| 170 | |
| 171 | self.gamma_1 = 1.0 |
| 172 | self.gamma_2 = 1.0 |
| 173 | if use_layerscale: |
| 174 | self.gamma_1 = nn.Parameter(layerscale_value * torch.ones((dim)), requires_grad=True) |
| 175 | self.gamma_2 = nn.Parameter(layerscale_value * torch.ones((dim)), requires_grad=True) |
| 176 | |
| 177 | self.H = None |
| 178 | self.W = None |
| 179 | |
| 180 | def forward(self, x): |
| 181 | H, W = self.H, self.W |
| 182 | B, L, C = x.shape |
| 183 | shortcut = x |
| 184 |