Untitled
unknown
plain_text
3 years ago
1.0 kB
3
Indexable
class TransformerBlock(nn.Module): # Vision Transformer https://arxiv.org/abs/2010.11929 def __init__(self, c1, c2, num_heads, num_layers): super().__init__() self.conv = None if c1 != c2: self.conv = Conv(c1, c2) self.linear = nn.Linear(c2, c2) # learnable position embedding self.tr = nn.Sequential(*(TransformerLayer(c2, num_heads) for _ in range(num_layers))) self.c2 = c2 self.to_patch = Conv(c1 * 5, c1) def forward(self, x): shifts = ((1, -1, 0, 0), (-1, 1, 0, 0), (0, 0, 1, -1), (0, 0, -1, 1)) shifted_x = list(map(lambda shift: F.pad(x, shift), shifts)) x_with_shifts = torch.cat((x, *shifted_x), dim = 1) x = self.to_patch(x_with_shifts) if self.conv is not None: x = self.conv(x) b, _, w, h = x.shape p = x.flatten(2).permute(2, 0, 1) return self.tr(p + self.linear(p)).permute(1, 2, 0).reshape(b, self.c2, w, h)
Editor is loading...