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
| 59 | |
| 60 | # Lightly adapted from ConvNext (https://github.com/facebookresearch/ConvNeXt) |
| 61 | class CXBlock(nn.Module): |
| 62 | r"""ConvNeXt Block. There are two equivalent implementations: |
| 63 | (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W) |
| 64 | (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back |
| 65 | We use (2) as we find it slightly faster in PyTorch |
| 66 | |
| 67 | Args: |
| 68 | dim (int): Number of input channels. |
| 69 | drop_path (float): Stochastic depth rate. Default: 0.0 |
| 70 | layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. |
| 71 | """ |
| 72 | |
| 73 | def __init__( |
| 74 | self, |
| 75 | dim, |
| 76 | kernel_size=7, |
| 77 | padding=3, |
| 78 | drop_path=0.0, |
| 79 | layer_scale_init_value=1e-6, |
| 80 | use_dwconv=True, |
| 81 | ): |
| 82 | super().__init__() |
| 83 | self.dwconv = nn.Conv2d( |
| 84 | dim, |
| 85 | dim, |
| 86 | kernel_size=kernel_size, |
| 87 | padding=padding, |
| 88 | groups=dim if use_dwconv else 1, |
| 89 | ) # depthwise conv |
| 90 | self.norm = LayerNorm2d(dim, eps=1e-6) |
| 91 | self.pwconv1 = nn.Linear(dim, 4 * dim) # pointwise/1x1 convs, implemented with linear layers |
| 92 | self.act = nn.GELU() |
| 93 | self.pwconv2 = nn.Linear(4 * dim, dim) |
| 94 | # NOTE: changed from gamma to weight |
| 95 | # https://github.com/huggingface/transformers/issues/29554 |
| 96 | self.weight = ( |
| 97 | nn.Parameter(layer_scale_init_value * torch.ones( |
| 98 | (dim)), requires_grad=True) if layer_scale_init_value > 0 else None) |
| 99 | self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() |
| 100 | |
| 101 | def forward(self, x): |
| 102 | input = x |
| 103 | x = self.dwconv(x) |
| 104 | x = self.norm(x) |
| 105 | x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C) |
| 106 | x = self.pwconv1(x) |
| 107 | x = self.act(x) |
| 108 | x = self.pwconv2(x) |
| 109 | if self.weight is not None: |
| 110 | x = self.weight * x |
| 111 | x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W) |
| 112 | |
| 113 | x = input + self.drop_path(x) |
| 114 | return x |
| 115 | |
| 116 | |
| 117 | class Fuser(nn.Module): |