Transformer-based OCR model not learning

EvanZ: <p>Would appreciate any thoughts or references to papers that might be relevant. I’m trying to build a simple OCR model (see example images from the MNIST words dataset) where I encode image features using <code>vgg</code> and feed those features into a Transformer. Everything is “working” in the sense that it seems to be going through the training without throwing any errors. But the results so far have been nonsense, although the training error does improve somewhat over the first few epochs. Here are the two main classes. I guess the most important question to address first is whether this architecture <em>should</em> work if done correctly (perhaps, I’m doing something incorrectly).</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/0/0/00e4f8b2af247d9b10d1a99c9f43819d16990a8c.png" alt="44" data-base62-sha1="7UzhIDWTMd6LHLPM5h1WpeuQos" width="115" height="47"><br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/c/8/c89336be302fd5254aa93e1c16cdd1dec2809af4.jpeg" alt="0" data-base62-sha1="sCmTH0NGrbS3oGFHbH2xfdjAAjq" width="22" height="21"><br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/9/2/92157b8b76ec9603fae3e95a2b149bbd0ae79441.jpeg" alt="1" data-base62-sha1="kQjJGMI52QBvzjENwRP5tjAVLvH" width="34" height="23"><br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/2/b/2b5e6ad0077c2cb3dd302855c26841ca9d4b9a74.jpeg" alt="2" data-base62-sha1="6bEON6wWgz7LXMJrP7P1Iv0ugmM" width="47" height="23"><br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/2/a/2abd2fa4aa867f068338fc2dca36c51a9af0b285.jpeg" alt="3" data-base62-sha1="665nKIMPWqod5mVEAh11OhhKhjD" width="79" height="37"></p>
<pre><code class="lang-auto">class FeatureExtractor(torch.nn.Module):
    def __init__(self):
        super(FeatureExtractor, self).__init__()
        vgg = torchvision.models.vgg16(pretrained=False)
        self.cnn = torch.nn.Sequential()
        vgg_modules = list(vgg.children())[:-1]
        for i in range(len(vgg_modules)):
            self.cnn.add_module(str(i), vgg_modules[i])
        self.cnn.add_module(str(i + 1), torch.nn.Sequential(
            torch.nn.Flatten(start_dim=2),
            torch.nn.Linear(in_features=49, out_features=96)
        ))
        self.pos_encoder = PositionalEncodingTrig(96, 0.1)

    def forward(self, x):
        x = self.cnn(x)
        x = x.reshape(512, -1, 96)
        x = self.pos_encoder(x)
        return x
</code></pre>
<pre><code class="lang-auto">class OCRModel(torch.nn.Module):
    def __init__(self):
        super(OCRModel, self).__init__()
        self.fe = FeatureExtractor()
        self.transformer = torch.nn.Transformer(d_model=96, nhead=8, num_encoder_layers=12,
                                                num_decoder_layers=12)
        self.linear = torch.nn.Linear(in_features=96, out_features=len(vocab))
        self.tgt_tok_emb = torch.nn.Embedding(num_embeddings=len(vocab), embedding_dim=96)
        self.pos_encoder = PositionalEncodingTrig(96)

    def forward(self, input, tgt):
        src = self.fe(input)
        tgt_embed = self.pos_encoder(self.tgt_tok_emb(tgt))
        tgt_mask = self.transformer.generate_square_subsequent_mask(tgt_embed.shape[0]).to(device)
        x = self.transformer(src, tgt_embed, tgt_mask=tgt_mask)
        x = self.linear(x)
        return x
</code></pre>