Multi-Head Convolutional Attention
| 128 | |
| 129 | |
| 130 | class MHCA(nn.Module): |
| 131 | """ |
| 132 | Multi-Head Convolutional Attention |
| 133 | """ |
| 134 | def __init__(self, out_channels, head_dim): |
| 135 | super(MHCA, self).__init__() |
| 136 | norm_layer = partial(nn.BatchNorm2d, eps=NORM_EPS) |
| 137 | self.group_conv3x3 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, |
| 138 | padding=1, groups=out_channels // head_dim, bias=False) |
| 139 | self.norm = norm_layer(out_channels) |
| 140 | self.act = nn.ReLU(inplace=True) |
| 141 | self.projection = nn.Conv2d(out_channels, out_channels, kernel_size=1, bias=False) |
| 142 | |
| 143 | def forward(self, x): |
| 144 | out = self.group_conv3x3(x) |
| 145 | out = self.norm(out) |
| 146 | out = self.act(out) |
| 147 | out = self.projection(out) |
| 148 | return out |
| 149 | |
| 150 | |
| 151 | class Mlp(nn.Module): |