Not running forward -- ViT reconstructed

username: xu555
cooked: <p>I am trying to remove the last fc layer of ViT and create my own fc layer.</p>
<p>I used this <a href="https://pytorch.org/vision/master/_modules/torchvision/models/vision_transformer.html" rel="noopener nofollow ugc">link</a>  to import <code>vit_b_16</code></p>
<pre><code class="lang-auto">class ViT_df(nn.Module):
    def __init__(self, num_classes=2, model_name='vit16', pretrained=True):
        super().__init__()
        if model_name == 'vit16':
            model = vit_b_16(pretrained=pretrained)

        else:
            print('wrong model name')

        layers = list(model.children())[:-1]
        fc_size = model.hidden_dim
        self.parent = nn.Sequential(*layers)
        self.fc = nn.Linear(in_features=fc_size, out_features=num_classes)

    def forward(self, image):
        x = self.parent(image)
        x = self.dropout(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)

        return x
</code></pre>
<p>However, it returns error:</p>
<blockquote>
<p>AssertionError: Expected (batch_size, seq_length, hidden_dim) got torch.Size([50, 768, 14, 14])</p>
</blockquote>
<p>batch_size=50, seq_length = 14*14, hidden_dim=768 and obviously this error is due to not running <code>def _process_input(self, x: torch.Tensor) -&gt; torch.Tensor:</code> function. (in this <a href="https://pytorch.org/vision/master/_modules/torchvision/models/vision_transformer.html" rel="noopener nofollow ugc">link</a>)</p>
<p>I debugged my code. My reconstructed ViT model runs <code>forward()</code> in <code>class Encoder(nn.Module):</code> before running <code>forward()</code> in <code>class VisionTransformer(nn.Module):</code>. (in this <a href="https://pytorch.org/vision/master/_modules/torchvision/models/vision_transformer.html" rel="noopener nofollow ugc">link</a>) It;s very weird for me. And I need your help.</p>
<p>Thank you!!!</p>

username: Matias_Vasquez
cooked: <p>What I assume is happening, is you are getting rid of the whole <code>self.heads</code> layer that comes after the <code>Encoder</code>, and not only the last <code>nn.Linear</code> that is inside of <code>self.heads</code>. You can print the list of layers that you are taking from the parent to be sure.</p>
<p>If this is the case, then you need to take only the <code>class token</code> before passing it through your custom <code>nn.Linear</code>.</p>
<pre><code class="lang-python"># Only the class token
x = x[:, 0]
</code></pre>
<p>This should work fine to fine-tune.</p>
<pre><code class="lang-python">class ViT_df(nn.Module):
    def __init__(self, num_classes=2, model_name='vit16', pretrained=True):
        super().__init__()
        if model_name == 'vit16':
            model = vit_b_16(pretrained=pretrained)

        else:
            print('wrong model name')

        layers = list(model.children())[:-1]
        fc_size = model.hidden_dim
        self.parent = nn.Sequential(*layers)
        self.fc = nn.Linear(in_features=fc_size, out_features=num_classes)

    def forward(self, image):
        x = self.parent(image)
        # Select ONLY the class token in the first column
        x = x[:, 0] 

        x = self.dropout(x)

        # I´m not sure you need this
        #x = x.view(x.size(0), -1)
        
        x = self.fc(x)

        return x
</code></pre>
<p>If you want to use the full MLP after the encoder, then you can either copy the full <code>self.heads</code> or make sure that you are only getting rid of the last layer inside <code>self.heads</code> and not the whole thing.</p>
<p>If you are getting rid of only the last layer and the problem is still there, please let me know.</p>
<p>Hope this helps.</p>
<p>P.S.: I am only assuming this, because in the link you posted, <code>self.heads</code> is defined as <code>nn.Sequential</code> with other layers as children. This means that the actual layer that you want to replace is a child of the child you are removing.</p>

username: xu555
cooked: <aside class="quote no-group" data-username="Matias_Vasquez" data-post="2" data-topic="147896">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/matias_vasquez/48/47847_2.png" class="avatar"> Matias_Vasquez:</div>
<blockquote>
<p><code>x = x[:, 0] </code></p>
</blockquote>
</aside>
<p>Thank you for your reply Matias, and I tried your suggestion. But I am afraid it is not my case.</p>
<ol>
<li>The error happens in <code>x = self.parent(image)</code>, so it does not help if <code>x = x[:, 0]</code> is added after it.</li>
<li>There is only one <code>nn.Linear</code> layer in <code>heads</code> if I don’t set representation_size which is my case.</li>
</ol>
<pre><code class="lang-auto">        if representation_size is None:
            heads_layers["head"] = nn.Linear(hidden_dim, num_classes)
        else:
            heads_layers["pre_logits"] = nn.Linear(hidden_dim, representation_size)
            heads_layers["act"] = nn.Tanh()
            heads_layers["head"] = nn.Linear(representation_size, num_classes)
</code></pre>

username: Matias_Vasquez
cooked: <p>Oh I see. As you mentioned, the image should have the shape B x Seq_len x Hidd_dim = 50 x (14*14 + class_token) x 768.</p>
<p>Could you print something like this to see in which exact layer the problem is?</p>
<pre><code class="lang-python">def forward(self, x):
    for layer in self.parent:
        print(layer, x.shape)
        x = layer(x)

    x = self.dropout(x)
    x = x.view(x.size(0), -1)
    x = self.fc(x)

    return x
</code></pre>

username: xu555
cooked: <p>Yes, here it is.</p>
<pre><code class="lang-auto">Conv2d(3, 768, kernel_size=(16, 16), stride=(16, 16)) torch.Size([50, 3, 224, 224])
Encoder(
  (dropout): Dropout(p=0.0, inplace=False)
  (layers): Sequential(
    (encoder_layer_0): EncoderBlock(
      (ln_1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (self_attention): MultiheadAttention(
        (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)
      )
      (dropout): Dropout(p=0.0, inplace=False)
      (ln_2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MLPBlock(
        (linear_1): Linear(in_features=768, out_features=3072, bias=True)
        (act): GELU()
        (dropout_1): Dropout(p=0.0, inplace=False)
        (linear_2): Linear(in_features=3072, out_features=768, bias=True)
        (dropout_2): Dropout(p=0.0, inplace=False)
      )
    )
    (encoder_layer_1): EncoderBlock(
      (ln_1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (self_attention): MultiheadAttention(
        (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)
      )
      (dropout): Dropout(p=0.0, inplace=False)
      (ln_2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MLPBlock(
        (linear_1): Linear(in_features=768, out_features=3072, bias=True)
        (act): GELU()
        (dropout_1): Dropout(p=0.0, inplace=False)
        (linear_2): Linear(in_features=3072, out_features=768, bias=True)
        (dropout_2): Dropout(p=0.0, inplace=False)
      )
    )
    (encoder_layer_2): EncoderBlock(
      (ln_1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (self_attention): MultiheadAttention(
        (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)
      )
      (dropout): Dropout(p=0.0, inplace=False)
      (ln_2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MLPBlock(
        (linear_1): Linear(in_features=768, out_features=3072, bias=True)
        (act): GELU()
        (dropout_1): Dropout(p=0.0, inplace=False)
        (linear_2): Linear(in_features=3072, out_features=768, bias=True)
        (dropout_2): Dropout(p=0.0, inplace=False)
      )
    )
    (encoder_layer_3): EncoderBlock(
      (ln_1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (self_attention): MultiheadAttention(
        (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)
      )
      (dropout): Dropout(p=0.0, inplace=False)
      (ln_2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MLPBlock(
        (linear_1): Linear(in_features=768, out_features=3072, bias=True)
        (act): GELU()
        (dropout_1): Dropout(p=0.0, inplace=False)
        (linear_2): Linear(in_features=3072, out_features=768, bias=True)
        (dropout_2): Dropout(p=0.0, inplace=False)
      )
    )
    (encoder_layer_4): EncoderBlock(
      (ln_1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (self_attention): MultiheadAttention(
        (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)
      )
      (dropout): Dropout(p=0.0, inplace=False)
      (ln_2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MLPBlock(
        (linear_1): Linear(in_features=768, out_features=3072, bias=True)
        (act): GELU()
        (dropout_1): Dropout(p=0.0, inplace=False)
        (linear_2): Linear(in_features=3072, out_features=768, bias=True)
        (dropout_2): Dropout(p=0.0, inplace=False)
      )
    )
    (encoder_layer_5): EncoderBlock(
      (ln_1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (self_attention): MultiheadAttention(
        (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)
      )
      (dropout): Dropout(p=0.0, inplace=False)
      (ln_2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MLPBlock(
        (linear_1): Linear(in_features=768, out_features=3072, bias=True)
        (act): GELU()
        (dropout_1): Dropout(p=0.0, inplace=False)
        (linear_2): Linear(in_features=3072, out_features=768, bias=True)
        (dropout_2): Dropout(p=0.0, inplace=False)
      )
    )
    (encoder_layer_6): EncoderBlock(
      (ln_1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (self_attention): MultiheadAttention(
        (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)
      )
      (dropout): Dropout(p=0.0, inplace=False)
      (ln_2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MLPBlock(
        (linear_1): Linear(in_features=768, out_features=3072, bias=True)
        (act): GELU()
        (dropout_1): Dropout(p=0.0, inplace=False)
        (linear_2): Linear(in_features=3072, out_features=768, bias=True)
        (dropout_2): Dropout(p=0.0, inplace=False)
      )
    )
    (encoder_layer_7): EncoderBlock(
      (ln_1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (self_attention): MultiheadAttention(
        (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)
      )
      (dropout): Dropout(p=0.0, inplace=False)
      (ln_2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MLPBlock(
        (linear_1): Linear(in_features=768, out_features=3072, bias=True)
        (act): GELU()
        (dropout_1): Dropout(p=0.0, inplace=False)
        (linear_2): Linear(in_features=3072, out_features=768, bias=True)
        (dropout_2): Dropout(p=0.0, inplace=False)
      )
    )
    (encoder_layer_8): EncoderBlock(
      (ln_1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (self_attention): MultiheadAttention(
        (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)
      )
      (dropout): Dropout(p=0.0, inplace=False)
      (ln_2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MLPBlock(
        (linear_1): Linear(in_features=768, out_features=3072, bias=True)
        (act): GELU()
        (dropout_1): Dropout(p=0.0, inplace=False)
        (linear_2): Linear(in_features=3072, out_features=768, bias=True)
        (dropout_2): Dropout(p=0.0, inplace=False)
      )
    )
    (encoder_layer_9): EncoderBlock(
      (ln_1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (self_attention): MultiheadAttention(
        (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)
      )
      (dropout): Dropout(p=0.0, inplace=False)
      (ln_2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MLPBlock(
        (linear_1): Linear(in_features=768, out_features=3072, bias=True)
        (act): GELU()
        (dropout_1): Dropout(p=0.0, inplace=False)
        (linear_2): Linear(in_features=3072, out_features=768, bias=True)
        (dropout_2): Dropout(p=0.0, inplace=False)
      )
    )
    (encoder_layer_10): EncoderBlock(
      (ln_1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (self_attention): MultiheadAttention(
        (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)
      )
      (dropout): Dropout(p=0.0, inplace=False)
      (ln_2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MLPBlock(
        (linear_1): Linear(in_features=768, out_features=3072, bias=True)
        (act): GELU()
        (dropout_1): Dropout(p=0.0, inplace=False)
        (linear_2): Linear(in_features=3072, out_features=768, bias=True)
        (dropout_2): Dropout(p=0.0, inplace=False)
      )
    )
    (encoder_layer_11): EncoderBlock(
      (ln_1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (self_attention): MultiheadAttention(
        (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)
      )
      (dropout): Dropout(p=0.0, inplace=False)
      (ln_2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MLPBlock(
        (linear_1): Linear(in_features=768, out_features=3072, bias=True)
        (act): GELU()
        (dropout_1): Dropout(p=0.0, inplace=False)
        (linear_2): Linear(in_features=3072, out_features=768, bias=True)
        (dropout_2): Dropout(p=0.0, inplace=False)
      )
    )
  )
  (ln): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
) torch.Size([50, 768, 14, 14])

</code></pre>

username: Matias_Vasquez
cooked: <p>You mentioned in your first post that you are not running <code>_process_input()</code>. I thought you meant that it was inside your Encoder or something like that. But you need to run this in order to have the expected shape.</p>
<p>Here is where the the patches are flattened.</p>
<p>Also I do not see the <code>class token</code> added.</p>
<p>Is there a reason you left the <code>_process_input</code> method out?</p>
<p>Edit: 1. In the implementation from your link, the positional embedding is added inside the encoder.<br>
2. You can access the <code>class_token</code> from the parent, since it is defined as <code>self.class_token = nn.Parameter(...)</code> in the <code>VisionTransformer</code> class.</p>

username: Matias_Vasquez
cooked: <p>The same issue is being discussed here.</p>
<aside class="quote" data-post="2" data-topic="148029">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/feature-extraction-in-torchvision-models-vit-b-16/148029/2">Feature extraction in torchvision.models.vit_b_16</a> <a class="badge-category__wrapper " href="/c/vision/5"><span data-category-id="5" data-drop-close="true" class="badge-category " title="Topics related to either pytorch/vision or vision research related topics"><span class="badge-category__name">vision</span></span></a>
  </div>
  <blockquote>
    Wrapping the submodules into an nn.Sequential container works for models which only use nn.Modules and no functional API calls and initialize them in a sequential manner. 
In your use case you would lose all functional calls used in the original forward method as seen <a href="https://github.com/pytorch/vision/blob/11bd2eaa6d6976129836b329b01d1300babddcc9/torchvision/models/vision_transformer.py#L258-L274" rel="noopener nofollow ugc">here</a>.
  </blockquote>
</aside>

<p>No functional API calls are done. In the comments I explain a bit of how this is done.</p>
<p>Please let me know if you still need help.</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn
import torchvision

img = torch.randn(1, 3, 224, 224)


model = torchvision.models.vit_b_16()
feature_extractor = nn.Sequential(*list(model.children())[:-1])

# This is supposed to be the PREPROCESS
# But it is not done correctly, since the reshaping and permutation is not done
# Only the concolution
conv = feature_extractor[0]  

# -&gt; print(conv(img).shape)
# -&gt; torch.Size([1, 768, 14, 14])
# This is not the desired output after preprocessing the image into
# flat patches. Also in the pytorch implementation, the class token
# and positional embedding are done extra on the forward method.

# This is the whole encoder sequence
encoder = feature_extractor[1]

# The MLP head at the end is gone, since you only selected the children until -1
# mlp = feature_extractor[2]

# This is how the model preprocess the image.
# The output shape is the one desired 
x = model._process_input(img)

# -&gt; print(x.shape)
# -&gt; torch.Size([1, 197, 768])
# This is Batch x N_Patches+Class_Token x C * H_patch * W_patch
# Meaning   1   x   14*14  +     1      x 3 * 16* 16   
       
# However, if you actually print the shape in here you only get 196 in dim=1
# This means that the class token in missing
# The positional_embedding is done inside the encoder, so I guess should be fine

# The next code is just copy paste from the forward method in the source code
# for the vit_b_16 from pytorch in order to get the 

n = x.shape[0]
# Expand the class token to the full batch
batch_class_token = model.class_token.expand(n, -1, -1)
x = torch.cat([batch_class_token, x], dim=1)
x = encoder(x)

# Classifier "token" as used by standard language architectures

x = x[:, 0]

# Here you can use your own nn.Linear to map to your number of classes
</code></pre>