r""" ConvNeXt Block. There are two equivalent implementations: (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W) (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back We use (2) as we find it
| 13 | from timm.models.registry import register_model |
| 14 | |
| 15 | class Block(nn.Module): |
| 16 | r""" ConvNeXt Block. There are two equivalent implementations: |
| 17 | (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W) |
| 18 | (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back |
| 19 | We use (2) as we find it slightly faster in PyTorch |
| 20 | |
| 21 | Args: |
| 22 | dim (int): Number of input channels. |
| 23 | drop_path (float): Stochastic depth rate. Default: 0.0 |
| 24 | layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. |
| 25 | """ |
| 26 | def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6): |
| 27 | super().__init__() |
| 28 | self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv |
| 29 | self.norm = LayerNorm(dim, eps=1e-6) |
| 30 | self.pwconv1 = nn.Linear(dim, 4 * dim) # pointwise/1x1 convs, implemented with linear layers |
| 31 | self.act = nn.GELU() |
| 32 | self.pwconv2 = nn.Linear(4 * dim, dim) |
| 33 | self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)), |
| 34 | requires_grad=True) if layer_scale_init_value > 0 else None |
| 35 | self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() |
| 36 | |
| 37 | def forward(self, x): |
| 38 | input = x |
| 39 | x = self.dwconv(x) |
| 40 | x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C) |
| 41 | x = self.norm(x) |
| 42 | x = self.pwconv1(x) |
| 43 | x = self.act(x) |
| 44 | x = self.pwconv2(x) |
| 45 | if self.gamma is not None: |
| 46 | x = self.gamma * x |
| 47 | x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W) |
| 48 | |
| 49 | x = input + self.drop_path(x) |
| 50 | return x |
| 51 | |
| 52 | class ConvNeXt(nn.Module): |
| 53 | r""" ConvNeXt |