| 14 | |
| 15 | |
| 16 | class Attention(nn.Module): |
| 17 | def __init__(self, in_ch, num_groups, D=3): |
| 18 | super(Attention, self).__init__() |
| 19 | assert in_ch % num_groups == 0 |
| 20 | if D == 3: |
| 21 | self.q = nn.Conv3d(in_ch, in_ch, 1) |
| 22 | self.k = nn.Conv3d(in_ch, in_ch, 1) |
| 23 | self.v = nn.Conv3d(in_ch, in_ch, 1) |
| 24 | |
| 25 | self.out = nn.Conv3d(in_ch, in_ch, 1) |
| 26 | elif D == 1: |
| 27 | self.q = nn.Conv1d(in_ch, in_ch, 1) |
| 28 | self.k = nn.Conv1d(in_ch, in_ch, 1) |
| 29 | self.v = nn.Conv1d(in_ch, in_ch, 1) |
| 30 | |
| 31 | self.out = nn.Conv1d(in_ch, in_ch, 1) |
| 32 | |
| 33 | self.norm = nn.GroupNorm(num_groups, in_ch) |
| 34 | self.nonlin = Swish() |
| 35 | |
| 36 | self.sm = nn.Softmax(-1) |
| 37 | |
| 38 | |
| 39 | def forward(self, x): |
| 40 | B, C = x.shape[:2] |
| 41 | h = x |
| 42 | |
| 43 | |
| 44 | |
| 45 | |
| 46 | q = self.q(h).reshape(B,C,-1) |
| 47 | k = self.k(h).reshape(B,C,-1) |
| 48 | v = self.v(h).reshape(B,C,-1) |
| 49 | |
| 50 | qk = torch.matmul(q.permute(0, 2, 1), k) #* (int(C) ** (-0.5)) |
| 51 | |
| 52 | w = self.sm(qk) |
| 53 | |
| 54 | h = torch.matmul(v, w.permute(0, 2, 1)).reshape(B,C,*x.shape[2:]) |
| 55 | |
| 56 | h = self.out(h) |
| 57 | |
| 58 | x = h + x |
| 59 | |
| 60 | x = self.nonlin(self.norm(x)) |
| 61 | |
| 62 | return x |
| 63 | |
| 64 | class PVConv(nn.Module): |
| 65 | def __init__(self, in_channels, out_channels, kernel_size, resolution, attention=False, |