Replacing Dense Layer with Transformer in Image Model?

username: Stanley_C
cooked: <p>I was hoping to replace one Dense Layer with a Transformer in image classification for hopefully better performance. I’m hoping to replace the classifier section after the feature extraction with a transformer block.  Does anyone have any tips/code to show how to do this? My current issue is that most transformer modes have target mask, but I’m guessing that won’t help when replacing a Dense Layer, and may even hurt it. I also think that since we are not using text, embedding is not needed. I’m guessing that positional encoding is useful, given that orders of features matter, but I’m not sure if a standard Positional Encoding to Transformer Block to Linear layer setup would work at all.<br>
Any tips?</p>

username: tom
cooked: <p>I would look at the vision transformer architecture (so e.g. DeiT-Tiny to name one).</p>
<p>They have an embedding layer, too, that splits the images into 16x16 blocks and applies a single convolution to many channels to then add a learned positional embedding.<br>
You could use the output of the convolutional part as the blocks (so skipping that conv in the vision embedding) and just the learned embeddings.<br>
The vision transformers also add a (learned) global point (the class token) to the patches. This is where they then hook the prediction head to at the end of the transformer. This class token originates from BERT and, combined with its attention, essentially replaces the pooling to go from patches to global information.</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/MathInf/toroidal/blob/bff09f725627e4629d464008dc7c5f9d6322ebad/toroidal/models.py#L5-L28">
  <header class="source">

      <a href="https://github.com/MathInf/toroidal/blob/bff09f725627e4629d464008dc7c5f9d6322ebad/toroidal/models.py#L5-L28" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/MathInf/toroidal/blob/bff09f725627e4629d464008dc7c5f9d6322ebad/toroidal/models.py#L5-L28" target="_blank" rel="noopener nofollow ugc">MathInf/toroidal/blob/bff09f725627e4629d464008dc7c5f9d6322ebad/toroidal/models.py#L5-L28</a></h4>



  <pre class="onebox">    <code class="lang-py">
      <ol class="start lines" start="5" style="counter-reset: li-counter 4 ;">
          <li>class VisionPatchEmbedding(torch.nn.Module):</li>
          <li>    def __init__(self, input_size, patch_size, n_channels):</li>
          <li>        super().__init__()</li>
          <li>        self.n_channels = n_channels</li>
          <li>        num_patches = (input_size ** 2) // (</li>
          <li>            patch_size ** 2</li>
          <li>        )  # only when this is divisible...</li>
          <li>        self.patch_to_vec = torch.nn.Conv2d(</li>
          <li>            3,</li>
          <li>            n_channels,</li>
          <li>            kernel_size=(patch_size, patch_size),</li>
          <li>            stride=(patch_size, patch_size),</li>
          <li>        )  # equivalent to reshape + linear</li>
          <li>        self.class_token = torch.nn.Parameter(</li>
          <li>            torch.randn(1, 1, n_channels)</li>
          <li>        )  # "global information"</li>
          <li>        self.pos_embedding = torch.nn.Parameter(</li>
          <li>            torch.randn(1, num_patches + 1, n_channels)</li>
          <li>        )</li>
          <li>
      </li>
</ol>
    </code>
  </pre>

  This file has been truncated. <a href="https://github.com/MathInf/toroidal/blob/bff09f725627e4629d464008dc7c5f9d6322ebad/toroidal/models.py#L5-L28" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>The file has a DeiTTiny class a bit further down.</p>
<p>Best regards</p>
<p>Thomas</p>