| 388 | |
| 389 | |
| 390 | class TransformerMLP(nn.Module): |
| 391 | |
| 392 | def __init__(self, channels, expansion, drop): |
| 393 | |
| 394 | super().__init__() |
| 395 | |
| 396 | self.dim1 = channels |
| 397 | self.dim2 = channels * expansion |
| 398 | self.chunk = nn.Sequential() |
| 399 | self.chunk.add_module('linear1', nn.Linear(self.dim1, self.dim2)) |
| 400 | self.chunk.add_module('act', nn.GELU()) |
| 401 | self.chunk.add_module('drop1', nn.Dropout(drop, inplace=True)) |
| 402 | self.chunk.add_module('linear2', nn.Linear(self.dim2, self.dim1)) |
| 403 | self.chunk.add_module('drop2', nn.Dropout(drop, inplace=True)) |
| 404 | |
| 405 | def forward(self, x): |
| 406 | |
| 407 | _, _, H, W = x.size() |
| 408 | x = einops.rearrange(x, 'b c h w -> b (h w) c') |
| 409 | x = self.chunk(x) |
| 410 | x = einops.rearrange(x, 'b (h w) c -> b c h w', h=H, w=W) |
| 411 | return x |
| 412 | |
| 413 | class LayerNormProxy(nn.Module): |
| 414 |
nothing calls this directly
no outgoing calls
no test coverage detected