| 425 | |
| 426 | |
| 427 | class TransformerMLPWithConv(nn.Module): |
| 428 | |
| 429 | def __init__(self, channels, expansion, drop): |
| 430 | |
| 431 | super().__init__() |
| 432 | |
| 433 | self.dim1 = channels |
| 434 | self.dim2 = channels * expansion |
| 435 | self.linear1 = nn.Sequential( |
| 436 | nn.Conv2d(self.dim1, self.dim2, 1, 1, 0), |
| 437 | # nn.GELU(), |
| 438 | # nn.BatchNorm2d(self.dim2, eps=1e-5) |
| 439 | ) |
| 440 | self.drop1 = nn.Dropout(drop, inplace=True) |
| 441 | self.act = nn.GELU() |
| 442 | # self.bn = nn.BatchNorm2d(self.dim2, eps=1e-5) |
| 443 | self.linear2 = nn.Sequential( |
| 444 | nn.Conv2d(self.dim2, self.dim1, 1, 1, 0), |
| 445 | # nn.BatchNorm2d(self.dim1, eps=1e-5) |
| 446 | ) |
| 447 | self.drop2 = nn.Dropout(drop, inplace=True) |
| 448 | self.dwc = nn.Conv2d(self.dim2, self.dim2, 3, 1, 1, groups=self.dim2) |
| 449 | |
| 450 | def forward(self, x): |
| 451 | |
| 452 | x = self.linear1(x) |
| 453 | x = self.drop1(x) |
| 454 | x = x + self.dwc(x) |
| 455 | x = self.act(x) |
| 456 | # x = self.bn(x) |
| 457 | x = self.linear2(x) |
| 458 | x = self.drop2(x) |
| 459 | |
| 460 | return x |
nothing calls this directly
no outgoing calls
no test coverage detected