Feature extraction in torchvision.models.vit_b_16

David_Torpey: <p>Hi</p>
<p>It’s easy enough to obtain output features from the CNNs in torchvision.models by doing this:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torchvision.models as models

model = models.resnet18()
feature_extractor = nn.Sequential(*list(model.children())[:-1])
output_features = feature_extractor(torch.randn(1, 3, 224, 224))
</code></pre>
<p>However, this does not work when I try it with torchvision.models.vit_b_16. It results in the following error:</p>
<pre><code class="lang-auto">AssertionError: Expected (batch_size, seq_length, hidden_dim) got torch.Size([1, 768, 14, 14])
</code></pre>
<p>Any help would be greatly appreciated</p>

ptrblck: <p>Wrapping the submodules into an <code>nn.Sequential</code> container works for models which only use <code>nn.Module</code>s and no functional API calls and initialize them in a sequential manner.<br>
In your use case you would lose all functional calls used in the original <code>forward</code> method as seen <a href="https://github.com/pytorch/vision/blob/11bd2eaa6d6976129836b329b01d1300babddcc9/torchvision/models/vision_transformer.py#L258-L274">here</a>.</p>

Matias_Vasquez: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> beat me to it <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=12" title=":smile:" class="emoji" alt=":smile:" loading="lazy" width="20" height="20"> .</p>
<p>Hope this helps.</p>
<pre><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>

David_Torpey: <p><a class="mention" href="/u/matias_vasquez">@Matias_Vasquez</a> <a class="mention" href="/u/ptrblck">@ptrblck</a><br>
Thanks so much. So the “representation” for that image in the ViT model is the first vector extracted using <code>x=x[:, 0]</code> line?</p>
<p>Thanks a lot!</p>

Matias_Vasquez: <p>Yes, the class token in the first row is where the classification is going to be made. However, this still has the size of the dimension of one patch.</p>
<p>This means that if each patch is 3*16*16 (CxHxW), then the size of x is</p>
<pre><code class="lang-python"># 14*14 patches
# 16*16 pixels in each patch

x.shape # -&gt; 197x768 = 14*14+1 x 16*16
x = x[:, 0]
x.shape # -&gt; 1x768
</code></pre>
<p>So afterwards comes the MLP after the encoder to reduce this to the number of classes you want. Some implementations use a couple of layers, but a single <code>nn. Linear(768, num_classes)</code> should be enough.</p>

David_Torpey: <p><a class="mention" href="/u/matias_vasquez">@Matias_Vasquez</a></p>
<p>Many thanks for your help. I was getting confused as to why one doesn’t take the mean of <code>x</code> when it has shape <code>197x768</code> to produce <code>1x768</code>. But I understand now why we instead take the first element to produce <code>1x768</code>.</p>
<p>Thanks!</p>

xu555: <p>Hello, do you have a solution for this problem? Thank you!</p>

ptrblck: <p><a class="mention" href="/u/matias_vasquez">@Matias_Vasquez</a> explained the proper approach really well. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"> Could you explain where you are stuck in this explanation?</p>

Denis_Mbey_Akola: <p>I am having a similar challenge. I am  trying to use vit_b_16 as my backbone for faster rcnn.</p>
<p>class VITWithFPN(torch.nn.Module):<br>
def <strong>init</strong>(self):<br>
super(VITWithFPN, self).<strong>init</strong>()<br>
# Get a  VIT backbone<br>
self.model =models.vit_b_16(pretrained=True,image_size=224)</p>
<pre><code>    #Extract the encoder      
    self.body = create_feature_extractor(
        self.model, return_nodes=['encoder'])
    
    inp = torch.randn(2, 3, 224, 224)
   
    with torch.no_grad():
        out = self.body(inp)
    in_channels_list = [o.shape[1] for o in out.values()]
    # Build FPN
    self.out_channels = self.model.hidden_dim
    self.fpn = torchvision.ops.FeaturePyramidNetwork(
        in_channels_list, out_channels=self.out_channels,
        extra_blocks=LastLevelMaxPool())

def _process_input(self, x: torch.Tensor) -&gt; torch.Tensor:
    n, c, h, w = x.shape
    print (n,c,h,w)
    print(x.shape)
    
    p = self.model.patch_size
    torch._assert(h == self.model.image_size, "Wrong image height!")
    torch._assert(w == self.model.image_size, "Wrong image width!")
    n_h = h // p
    n_w = w // p

    # (n, c, h, w) -&gt; (n, hidden_dim, n_h, n_w)
    x = self.model.conv_proj(x)
    # (n, hidden_dim, n_h, n_w) -&gt; (n, hidden_dim, (n_h * n_w))
    x = x.reshape(n, self.model.hidden_dim, n_h * n_w)

    # (n, hidden_dim, (n_h * n_w)) -&gt; (n, (n_h * n_w), hidden_dim)
    # The self attention layer expects inputs in the format (N, S, E)
    # where S is the source sequence length, N is the batch size, E is the
    # embedding dimension
    x = x.permute(0, 2, 1)

    return x

def forward(self, x):
    print("Image is --&gt;",x.size())
    x = self._process_input(x)
    n = x.shape[0]
    # Expand the class token to the full batch
    batch_class_token = self.model.class_token.expand(n, -1, -1)
    x = torch.cat([batch_class_token, x], dim=1)
    x = self.body(x)
    x = self.fpn(x)
    return x
</code></pre>
<p><span class="hashtag">#Build</span> model<br>
model = FasterRCNN(VITWithFPN(), num_classes=2)</p>
<p>I resized the images in my dataloader to 224*224. However,before the image gets through the forward method, the image tensors gets transformed  and have dimension torch.Size([8, 3, 800, 800])<br>
instead of torch.Size([3, 224, 224]). I would appreciate your help.<br>
Thanks in advance <a class="mention" href="/u/ptrblck">@ptrblck</a> <a class="mention" href="/u/matias_vasquez">@Matias_Vasquez</a></p>

Matias_Vasquez: <p>Hi,</p>
<p>sorry for the late reply. Could you where the shape changes? Or where this error occurs?</p>

Denis_Mbey_Akola: <p>The shape changes before the first line of code in forward function</p>

Matias_Vasquez: <p>Ok,</p>
<p>since you are printing the shape of <code>x</code> in the first line of your forward function, this means that the ViT is <strong>NOT</strong> changing the size of your image, but whatever you are doing before.</p>
<p>I assume that the change in size is happening in your FasterRCNN <strong>before</strong> you feed it to the ViT.</p>
<aside class="quote no-group" data-username="Denis_Mbey_Akola" data-post="9" 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/denis_mbey_akola/48/48634_2.png" class="avatar"> Denis_Mbey_Akola:</div>
<blockquote>
<p><span class="hashtag-raw">#Build</span> model<br>
model = FasterRCNN(VITWithFPN(), num_classes=2)</p>
</blockquote>
</aside>
<p>If you post the code on how you create this FasterRCNN, then I maybe I can help you a little more.</p>

Denis_Mbey_Akola: <p>import os<br>
import configparser<br>
import torchvision<br>
import csv<br>
import os<br>
import os.path as osp<br>
import pickle<br>
from PIL import Image<br>
import numpy as np<br>
import scipy<br>
import torch<br>
import matplotlib.pyplot as plt<br>
from typing import overload</p>
<p>import torchvision.transforms as T</p>
<p>import utils<br>
from dataloader import MOT17ObjDetect<br>
from engine import train_one_epoch, evaluate<br>
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor<br>
from torchvision.models.detection.rpn import AnchorGenerator<br>
import torchvision.models as models<br>
from torchvision.models.feature_extraction import get_graph_node_names, create_feature_extractor<br>
import torchvision.models as models<br>
import torchvision<br>
from torchvision.models.detection.backbone_utils import LastLevelMaxPool<br>
from torchvision.models.detection import FasterRCNN<br>
import torch.nn as nn</p>
<p>from dataloader import MOT17ObjDetect<br>
import torchvision<br>
OUTPUT_DIR=“outputs”<br>
LOG_RESULTS =“inference”<br>
PREDICTED_IMAGES =“images_with_boxes”<br>
import random<br>
from torchvision.models.vision_transformer import vit_b_16</p>
<p>def plot(img, boxes):<br>
x=random.randint(100000, 100000000)<br>
fig, ax = plt.subplots(1, dpi=96)</p>
<p>img = img.mul(255).permute(1, 2, 0).cpu().byte().numpy()<br>
width, height, _ = img.shape</p>
<p>ax.imshow(img, cmap=‘gray’)<br>
fig.set_size_inches(width / 80, height / 80)</p>
<p>for box in boxes:<br>
rect = plt.Rectangle(<br>
(box[0], box[1]),<br>
box[2] - box[0],<br>
box[3] - box[1],<br>
fill=False,<br>
linewidth=1.0)<br>
ax.add_patch(rect)</p>
<p>plt.axis(‘off’)<br>
plt.savefig("./images_with_boxes/"+str(x)+".png")</p>
<p>class VITWithFPN(torch.nn.Module):<br>
def <strong>init</strong>(self):<br>
super(VITWithFPN, self).<strong>init</strong>()<br>
# Get a  VIT backbone<br>
self.model =models.vit_b_16(pretrained=True,image_size=224)</p>
<pre><code>    #Extract the encoder      
    self.body = create_feature_extractor(
        self.model, return_nodes=['encoder'])
    
    inp = torch.randn(2, 3, 224, 224)
   
    with torch.no_grad():
        out = self.body(inp)
    in_channels_list = [o.shape[1] for o in out.values()]
    # Build FPN
    self.out_channels = self.model.hidden_dim
    self.fpn = torchvision.ops.FeaturePyramidNetwork(
        in_channels_list, out_channels=self.out_channels,
        extra_blocks=LastLevelMaxPool())

def _process_input(self, x: torch.Tensor) -&gt; torch.Tensor:
    n, c, h, w = x.shape
    print (n,c,h,w)
    print(x.shape)
    
    p = self.model.patch_size
    torch._assert(h == self.model.image_size, "Wrong image height!")
    torch._assert(w == self.model.image_size, "Wrong image width!")
    n_h = h // p
    n_w = w // p

    # (n, c, h, w) -&gt; (n, hidden_dim, n_h, n_w)
    x = self.model.conv_proj(x)
    # (n, hidden_dim, n_h, n_w) -&gt; (n, hidden_dim, (n_h * n_w))
    x = x.reshape(n, self.model.hidden_dim, n_h * n_w)

    # (n, hidden_dim, (n_h * n_w)) -&gt; (n, (n_h * n_w), hidden_dim)
    # The self attention layer expects inputs in the format (N, S, E)
    # where S is the source sequence length, N is the batch size, E is the
    # embedding dimension
    x = x.permute(0, 2, 1)

    return x

def forward(self, x):
    print("Image is --&gt;",x.size())
    x = self._process_input(x)
    n = x.shape[0]
    # Expand the class token to the full batch
    batch_class_token = self.model.class_token.expand(n, -1, -1)
    x = torch.cat([batch_class_token, x], dim=1)
    x = self.body(x)
    x = self.fpn(x)
    return x
</code></pre>
<p><span class="hashtag">#Build</span> model<br>
model = FasterRCNN(VITWithFPN(), num_classes=2)</p>
<p>def get_transform(train):<br>
transforms = []<br>
# converts the image, a PIL image, into a PyTorch Tensor<br>
transforms.append(T.ToTensor())<br>
if train:<br>
# during training, randomly flip the training images<br>
# and ground-truth for data augmentation<br>
transforms.append(T.RandomHorizontalFlip(0.5))<br>
transforms.append(T.Resize((224,224)))<br>
return T.Compose(transforms)</p>
<h1>
<a name="use-our-dataset-and-defined-transformations-1" class="anchor" href="#use-our-dataset-and-defined-transformations-1"></a>use our dataset and defined transformations</h1>
<p>dataset = MOT17ObjDetect(’./data/MOT17Det/train’, get_transform(train=True))<br>
dataset_no_random = MOT17ObjDetect(’./data/MOT17Det/train’, get_transform(train=False))<br>
dataset_test = MOT17ObjDetect(’./data/MOT17Det/test’, get_transform(train=False))</p>
<h1>
<a name="split-the-dataset-in-train-and-test-set-2" class="anchor" href="#split-the-dataset-in-train-and-test-set-2"></a>split the dataset in train and test set</h1>
<p>torch.manual_seed(1)</p>
<h1>
<a name="define-training-and-validation-data-loaders-3" class="anchor" href="#define-training-and-validation-data-loaders-3"></a>define training and validation data loaders</h1>
<p>data_loader = torch.utils.data.DataLoader(<br>
dataset, batch_size=8, shuffle=True, num_workers=4,<br>
collate_fn=utils.collate_fn)<br>
data_loader_no_random = torch.utils.data.DataLoader(<br>
dataset_no_random, batch_size=8, shuffle=False, num_workers=4,<br>
collate_fn=utils.collate_fn)</p>
<p>data_loader_test = torch.utils.data.DataLoader(<br>
dataset_test, batch_size=2, shuffle=False, num_workers=4,<br>
collate_fn=utils.collate_fn)</p>
<p>device = torch.device(‘cuda’) if torch.cuda.is_available() else torch.device(‘cpu’)</p>
<h1>
<a name="get-the-model-using-our-helper-function-4" class="anchor" href="#get-the-model-using-our-helper-function-4"></a>get the model using our helper function</h1>
<h1>
<a name="model-get_detection_modeldatasetnum_classes-5" class="anchor" href="#model-get_detection_modeldatasetnum_classes-5"></a>model = get_detection_model(dataset.num_classes)</h1>
<h1>
<a name="move-model-to-the-right-device-6" class="anchor" href="#move-model-to-the-right-device-6"></a>move model to the right device</h1>
<p>model.to(device)</p>
<h1>
<a name="construct-an-optimizer-7" class="anchor" href="#construct-an-optimizer-7"></a>construct an optimizer</h1>
<p>params = [p for p in model.parameters() if p.requires_grad]<br>
optimizer = torch.optim.SGD(params, lr=0.00125,<br>
momentum=0.9, weight_decay=0.0005)</p>
<h1>
<a name="and-a-learning-rate-scheduler-which-decreases-the-learning-rate-by-8" class="anchor" href="#and-a-learning-rate-scheduler-which-decreases-the-learning-rate-by-8"></a>and a learning rate scheduler which decreases the learning rate by</h1>
<h1>
<a name="h-10x-every-3-epochs-9" class="anchor" href="#h-10x-every-3-epochs-9"></a>10x every 3 epochs</h1>
<p>lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,<br>
step_size=10,<br>
gamma=0.1)<br>
os.makedirs(OUTPUT_DIR,exist_ok=True)<br>
os.makedirs(LOG_RESULTS,exist_ok=True)<br>
def evaluate_and_write_result_files(model, data_loader):<br>
model.eval()<br>
results = {}<br>
for imgs, targets in data_loader:<br>
imgs = [img.to(device) for img in imgs]</p>
<pre><code>with torch.no_grad():
    preds = model(imgs)

for pred, target in zip(preds, targets):
    results[target['image_id'].item()] = {'boxes': pred['boxes'].cpu(),
                                          'scores': pred['scores'].cpu()}
</code></pre>
<p>data_loader.dataset.print_eval(results)<br>
data_loader.dataset.write_results_files(results, OUTPUT_DIR+"/")</p>
<p>num_epochs = 27</p>
<p>for epoch in range(1, num_epochs + 1):<br>
train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=300)<br>
# update the learning rate<br>
# lr_scheduler.step()<br>
# evaluate on the test dataset<br>
if epoch % 3 == 0:<br>
evaluate_and_write_result_files(model, data_loader_no_random)<br>
torch.save(model.state_dict(), f"{OUTPUT_DIR}/model_epoch_{epoch}.model")</p>
<pre><code># pick one image from the test set
</code></pre>
<p>os.makedirs(PREDICTED_IMAGES,exist_ok=True)<br>
dataset = MOT17ObjDetect(’./data/MOT17Det/train’, get_transform(train=False))<br>
data_loader = torch.utils.data.DataLoader(<br>
dataset, batch_size=20, shuffle=True, num_workers=4,<br>
collate_fn=utils.collate_fn)</p>
<p>print("---------------Done training------------------")<br>
for imgs, target in data_loader:</p>
<pre><code># put the model in evaluation mode
model.eval()



with torch.no_grad():
    for k in range(0,20):
        prediction = model([imgs[k].to(device)])[0]
        

        plot(imgs[k], prediction['boxes'].cpu())
</code></pre>
<p><a class="mention" href="/u/matias_vasquez">@Matias_Vasquez</a> ,here is the complete code file</p>

Matias_Vasquez: <p>Since you are using this ↓ to declare your FasterRCNN, you can use one of the Arguments to make sure that the image has the right size before being fed to your backbone.</p>
<aside class="quote no-group" data-username="Denis_Mbey_Akola" data-post="13" 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/denis_mbey_akola/48/48634_2.png" class="avatar"> Denis_Mbey_Akola:</div>
<blockquote>
<p>from torchvision.models.detection import FasterRCNN</p>
</blockquote>
</aside>
<p><a href="https://pytorch.org/vision/stable/_modules/torchvision/models/detection/faster_rcnn.html" rel="noopener nofollow ugc">Here</a> is the documentation so you can see all of the arguments that you can use for this object. The <code>default</code> <code>min_size</code> is <code>800</code>. If you define a <code>max_size=224</code> this should be solved.</p>
<pre data-code-wrap="python"><code class="lang-python">model = FasterRCNN(VITWithFPN(), num_classes=2, max_size=224)
</code></pre>
<p>However I think there are still some issues with the Feature Pyramid.<br>
Also you are not selecting the <code>class token</code> as explained in the previous posts, but if that is what you want then you can try it out.</p>
<p>Hope this helps <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=12" title=":smile:" class="emoji" alt=":smile:" loading="lazy" width="20" height="20"></p>

Denis_Mbey_Akola: <p>Yes,there is still an issue. Please could you point out the error with my feature pyramid network?</p>

Matias_Vasquez: <p>Sorry but I do not quite understand HOW you want to use the FPN with the ViT.</p>
<p>If you declare your <code>body</code> like this ↓ and you print it, you will see that you do not have only the encoder. You have the entire ViT: <code>preprocess</code>, <code>encoder</code> and <code>head</code> (<code>LayerNorm</code>).</p>
<aside class="quote no-group" data-username="Denis_Mbey_Akola" data-post="13" 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/denis_mbey_akola/48/48634_2.png" class="avatar"> Denis_Mbey_Akola:</div>
<blockquote>
<pre><code class="lang-auto">self.body = create_feature_extractor(
        self.model, return_nodes=['encoder'])
</code></pre>
</blockquote>
</aside>
<p>This means that when you do something like this ↓ the input will pass through the entire ViT and return when you finish with the encoder (before the head). This means you do not have to do the <code>preprocess</code> yourself when doing the forward.</p>
<aside class="quote no-group" data-username="Denis_Mbey_Akola" data-post="13" 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/denis_mbey_akola/48/48634_2.png" class="avatar"> Denis_Mbey_Akola:</div>
<blockquote>
<pre><code class="lang-auto">with torch.no_grad():
        out = self.body(inp)
</code></pre>
</blockquote>
</aside>
<p>As mentioned, the output for this is the output of the encoder. Which has shape <code>torch.Size([1, 197, 768])</code>. As mentioned in previous posts, here is where you normally take the <code>class token</code> and pass it through a <code>MLP</code> to get the classification.</p>
<p>This is only <strong>ONE</strong> output tensor, and not the entire <code>encoder</code> sequence. So for your FPN you are only using this output to define your <code>in_channels_list</code>.</p>
<p>So, it really depends how you want to use your FPN and if it really is the way to go, instead of a normal <code>MLP</code>. But I do not know.</p>

Denis_Mbey_Akola: <p>Hello <a class="mention" href="/u/matias_vasquez">@Matias_Vasquez</a> ,thanks so much. I am really new to using ViT as a backbone. Please how will I adapt the MLP to use it in my FPN? Thanks in advance</p>

diegoaichele: <p>You can use another two way to the same:<br>
1°:  <a href="https://pytorch.org/vision/stable/feature_extraction.html" rel="noopener nofollow ugc">Feature extraction for model inspection — Torchvision 0.12 documentation (pytorch.org)</a></p>
<pre><code class="lang-auto">import torch
import torchvision.models as models
from torchvision.models.feature_extraction import create_feature_extractor, get_graph_node_names
# print(get_graph_node_names(model) )

model = models.vit_b_16()
x = torch.randn(1, 3, 224, 224)
output_feature = create_feature_extractor(model, return_nodes = { "encoder.ln":"getitem_5"} )
output_feature(x)
</code></pre>
<p>2°: <a href="https://pytorch.org/docs/stable/generated/torch.nn.modules.module.register_module_forward_hook.html" rel="noopener nofollow ugc">torch.nn.modules.module.register_module_forward_hook — PyTorch 1.11.0 documentation</a></p>
<pre><code class="lang-auto">import torch
import torchvision.models as models
model = models.vit_b_16()
def print_middle_layer(model,input,output):
    print("Print Output:", output)

model.encoder.ln.register_forward_hook( print_middle_layer ) 
# nn.Sequential(*list(model.children())[:-1]) == model.encoder.ln
x = torch.randn(1, 3, 224, 224)
model(x)
</code></pre>

Denis_Mbey_Akola: <p>Thanks,basically,what I am trying to do is to use vit_b_16() as my feature extractor for Faster RCNN instead of Resnet architecture (ResNet50,ResNet101,and ResNet152). I just runned your code <a class="mention" href="/u/diegoaichele">@diegoaichele</a>  and I am getting some errors.<br>
I want to use vit_b_16() as my backbone for faster rcnn. Maybe if you could help me do that.</p>

nicholasprimiano: <p>Im having the same issue. Were you ever able to get it working?</p>