| 73 | |
| 74 | |
| 75 | class PAM(nn.Module): |
| 76 | def __init__(self, in_channels): |
| 77 | super(PAM, self).__init__() |
| 78 | self.query_conv = nn.Conv2d(in_channels=in_channels, out_channels=in_channels // 4, kernel_size=1) |
| 79 | self.key_conv = nn.Conv2d(in_channels=in_channels, out_channels=in_channels // 4, kernel_size=1) |
| 80 | self.value_conv = nn.Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=1) |
| 81 | self.gamma = nn.Parameter(torch.zeros(1)) |
| 82 | self.softmax = nn.Softmax(dim=-1) |
| 83 | |
| 84 | def forward(self, x): |
| 85 | m_batchsize, C, height, width = x.size() |
| 86 | proj_query = self.query_conv(x).view(m_batchsize, -1, width * height).permute(0, 2, 1) |
| 87 | proj_key = self.key_conv(x).view(m_batchsize, -1, width * height) |
| 88 | energy = torch.bmm(proj_query, proj_key) |
| 89 | attention = self.softmax(energy) |
| 90 | proj_value = self.value_conv(x).view(m_batchsize, -1, width * height) |
| 91 | out = torch.bmm(proj_value, attention.permute(0, 2, 1)) |
| 92 | out = out.view(m_batchsize, C, height, width) |
| 93 | out = self.gamma * out + x |
| 94 | |
| 95 | return out |
| 96 | |
| 97 | |
| 98 | """ |