| 177 | |
| 178 | class ConvNeXtBlock(nn.Module): |
| 179 | def __init__(self, dim, drop_path=0.0, layer_scale_init_value=1e-6): |
| 180 | super().__init__() |
| 181 | self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) |
| 182 | self.norm = nn.LayerNorm(dim, eps=1e-6) |
| 183 | self.pwconv1 = nn.Linear(dim, 4 * dim) |
| 184 | self.act = nn.GELU() |
| 185 | self.pwconv2 = nn.Linear(4 * dim, dim) |
| 186 | self.gamma = nn.Parameter(layer_scale_init_value * torch.ones(dim)) if layer_scale_init_value > 0 else None |
| 187 | |
| 188 | from .modules import DropPath |
| 189 | self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() |
| 190 | |
| 191 | def forward(self, x): |
| 192 | identity = x |