| 13 | |
| 14 | |
| 15 | class Obj_Attn_Block(Module): |
| 16 | def __init__(self, in_dim, compress): |
| 17 | super(Obj_Attn_Block, self).__init__() |
| 18 | channel_in = in_dim//int(2*compress) |
| 19 | self.value_conv = Conv2d(in_channels=in_dim, out_channels=channel_in, kernel_size=1) |
| 20 | self.query_conv = Conv2d(in_channels=channel_in, out_channels=channel_in, kernel_size=1) |
| 21 | self.key_conv = Conv2d(in_channels=channel_in, out_channels=channel_in, kernel_size=1) |
| 22 | self.gamma = Parameter(torch.zeros(1), requires_grad=True) |
| 23 | self.softmax = Softmax(dim=-1) |
| 24 | |
| 25 | for layer in [self.value_conv, self.query_conv, self.key_conv]: |
| 26 | weight_init(layer) |
| 27 | |
| 28 | def forward(self, x): |
| 29 | m_batchsize, C, length, _ = x.size() |
| 30 | proj_value = self.value_conv(x).view(m_batchsize, -1, length) |
| 31 | x = proj_value.view(m_batchsize, -1, length, 1) |
| 32 | proj_query = self.query_conv(x).view(m_batchsize, -1, length).permute(0, 2, 1) |
| 33 | proj_key = self.key_conv(x).view(m_batchsize, -1, length) |
| 34 | energy = torch.bmm(proj_query, proj_key) |
| 35 | attention = self.softmax(energy) |
| 36 | x = torch.bmm(proj_value, attention.permute(0, 2, 1)) |
| 37 | x = torch.cat((self.gamma*x, proj_value), dim=1).view(m_batchsize, -1, length, 1) |
| 38 | return x |
| 39 | |
| 40 | |
| 41 | class OAM_GRAM(Module): |