(self, in_channels)
| 162 | |
| 163 | class SelfAttention2D(nn.Module): |
| 164 | def __init__(self, in_channels): |
| 165 | super().__init__() |
| 166 | self.query = nn.Conv2d(in_channels, in_channels // 8, 1) |
| 167 | self.key = nn.Conv2d(in_channels, in_channels // 8, 1) |
| 168 | self.value = nn.Conv2d(in_channels, in_channels, 1) |
| 169 | self.gamma = nn.Parameter(torch.zeros(1)) |
| 170 | |
| 171 | def forward(self, x): |
| 172 | B, C, H, W = x.size() |