Locally Connected Layer. Args: in_channels (int): the in channel of the features. out_channels (int): the out channel of the features. output_size (List[int]): the output size of the features. kernel_size (int): the
| 14 | |
| 15 | |
| 16 | class LocallyConnected2d(nn.Module): |
| 17 | """Locally Connected Layer. |
| 18 | |
| 19 | Args: |
| 20 | in_channels (int): |
| 21 | the in channel of the features. |
| 22 | out_channels (int): |
| 23 | the out channel of the features. |
| 24 | output_size (List[int]): |
| 25 | the output size of the features. |
| 26 | kernel_size (int): |
| 27 | the size of the kernel. |
| 28 | stride (int): |
| 29 | the stride of the kernel. |
| 30 | Returns: |
| 31 | attended_features (torch.Tensor): |
| 32 | attended feature maps |
| 33 | """ |
| 34 | def __init__(self, |
| 35 | in_channels, |
| 36 | out_channels, |
| 37 | output_size, |
| 38 | kernel_size, |
| 39 | stride, |
| 40 | bias=False): |
| 41 | super(LocallyConnected2d, self).__init__() |
| 42 | output_size = _pair(output_size) |
| 43 | self.weight = nn.Parameter( |
| 44 | torch.randn(1, out_channels, in_channels, output_size[0], |
| 45 | output_size[1], kernel_size**2), |
| 46 | requires_grad=True, |
| 47 | ) |
| 48 | if bias: |
| 49 | self.bias = nn.Parameter(torch.randn(1, out_channels, |
| 50 | output_size[0], |
| 51 | output_size[1]), |
| 52 | requires_grad=True) |
| 53 | else: |
| 54 | self.register_parameter('bias', None) |
| 55 | self.kernel_size = _pair(kernel_size) |
| 56 | self.stride = _pair(stride) |
| 57 | |
| 58 | def forward(self, x): |
| 59 | _, c, h, w = x.size() |
| 60 | kh, kw = self.kernel_size |
| 61 | dh, dw = self.stride |
| 62 | x = x.unfold(2, kh, dh).unfold(3, kw, dw) |
| 63 | x = x.contiguous().view(*x.size()[:-2], -1) |
| 64 | # Sum in in_channel and kernel_size dims |
| 65 | out = (x.unsqueeze(1) * self.weight).sum([2, -1]) |
| 66 | if self.bias is not None: |
| 67 | out += self.bias |
| 68 | return out |
| 69 | |
| 70 | |
| 71 | class KeypointAttention(nn.Module): |