sneedium/captcha_pixelplanet

1import torch2import torch.nn as nn3from .transformer import PositionalEncoding4 5class Attention(nn.Module):6    def __init__(self, in_channels=512, max_length=25, n_feature=256):7        super().__init__()8        self.max_length = max_length9 10        self.f0_embedding = nn.Embedding(max_length, in_channels)11        self.w0 = nn.Linear(max_length, n_feature)12        self.wv = nn.Linear(in_channels, in_channels)13        self.we = nn.Linear(in_channels, max_length)14 15        self.active = nn.Tanh()16        self.softmax = nn.Softmax(dim=2)17 18    def forward(self, enc_output):19        enc_output = enc_output.permute(0, 2, 3, 1).flatten(1, 2)20        reading_order = torch.arange(self.max_length, dtype=torch.long, device=enc_output.device)21        reading_order = reading_order.unsqueeze(0).expand(enc_output.size(0), -1)  # (S,) -> (B, S)22        reading_order_embed = self.f0_embedding(reading_order)  # b,25,51223 24        t = self.w0(reading_order_embed.permute(0, 2, 1))  # b,512,25625        t = self.active(t.permute(0, 2, 1) + self.wv(enc_output))  # b,256,51226 27        attn = self.we(t)  # b,256,2528        attn = self.softmax(attn.permute(0, 2, 1))  # b,25,25629        g_output = torch.bmm(attn, enc_output)  # b,25,51230        return g_output, attn.view(*attn.shape[:2], 8, 32)31 32 33def encoder_layer(in_c, out_c, k=3, s=2, p=1):34    return nn.Sequential(nn.Conv2d(in_c, out_c, k, s, p),35                         nn.BatchNorm2d(out_c),36                         nn.ReLU(True))37 38def decoder_layer(in_c, out_c, k=3, s=1, p=1, mode='nearest', scale_factor=None, size=None):39    align_corners = None if mode=='nearest' else True40    return nn.Sequential(nn.Upsample(size=size, scale_factor=scale_factor, 41                                     mode=mode, align_corners=align_corners),42                         nn.Conv2d(in_c, out_c, k, s, p),43                         nn.BatchNorm2d(out_c),44                         nn.ReLU(True))45 46 47class PositionAttention(nn.Module):48    def __init__(self, max_length, in_channels=512, num_channels=64, 49                 h=8, w=32, mode='nearest', **kwargs):50        super().__init__()51        self.max_length = max_length52        self.k_encoder = nn.Sequential(53            encoder_layer(in_channels, num_channels, s=(1, 2)),54            encoder_layer(num_channels, num_channels, s=(2, 2)),55            encoder_layer(num_channels, num_channels, s=(2, 2)),56            encoder_layer(num_channels, num_channels, s=(2, 2))57        )58        self.k_decoder = nn.Sequential(59            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),60            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),61            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),62            decoder_layer(num_channels, in_channels, size=(h, w), mode=mode)63        )64 65        self.pos_encoder = PositionalEncoding(in_channels, dropout=0, max_len=max_length)66        self.project = nn.Linear(in_channels, in_channels)67 68    def forward(self, x):69        N, E, H, W = x.size()70        k, v = x, x  # (N, E, H, W)71 72        # calculate key vector73        features = []74        for i in range(0, len(self.k_encoder)):75            k = self.k_encoder[i](k)76            features.append(k)77        for i in range(0, len(self.k_decoder) - 1):78            k = self.k_decoder[i](k)79            k = k + features[len(self.k_decoder) - 2 - i]80        k = self.k_decoder[-1](k)81 82        # calculate query vector83        # TODO q=f(q,k)84        zeros = x.new_zeros((self.max_length, N, E))  # (T, N, E)85        q = self.pos_encoder(zeros)  # (T, N, E)86        q = q.permute(1, 0, 2)  # (N, T, E)87        q = self.project(q)  # (N, T, E)88        89        # calculate attention90        attn_scores = torch.bmm(q, k.flatten(2, 3))  # (N, T, (H*W))91        attn_scores = attn_scores / (E ** 0.5)92        attn_scores = torch.softmax(attn_scores, dim=-1)93 94        v = v.permute(0, 2, 3, 1).view(N, -1, E)  # (N, (H*W), E)95        attn_vecs = torch.bmm(attn_scores, v)  # (N, T, E)96 97        return attn_vecs, attn_scores.view(N, -1, H, W)