TypeError: avg_pool2d(): argument 'kernel_size' (position 2) must be tuple of ints, not tuple

username: whatevername
cooked: <p>Hi. I’m trying to implement a pre-trained Vision Transformer and perform gem pooling but I get this error.</p>
<p>The code is:</p>
<pre><code class="lang-auto">class GeM(nn.Module):
    
    def __init__(self, p=3, eps=1e-6):
        super(GeM,self).__init__()
        self.p = Parameter(torch.ones(1)*p)
        self.eps = eps

    def forward(self, x):
        return gem(x, p=self.p, eps=self.eps)
        
    def __repr__(self):
        return self.__class__.__name__ + '(' + 'p=' + '{:.4f}'.format(self.p.data.tolist()[0]) + ', ' + 'eps=' + str(self.eps) + ')'

def gem(x, p=3, eps=1e-6):

    return F.avg_pool2d(x.clamp(min=eps).pow(p), (x.size(-2), x.size(-1))).pow(1./p)

class Embedder(nn.Module):

    def __init__(self,architecture,gem_p=3,pretrained_flag = True):

        super(Embedder, self).__init__()

        #load the base model from PyTorch's pretrained models (imagenet pretrained)
        network = getattr(models,architecture)(pretrained=pretrained_flag)
        
        if architecture.startswith('vit'):
            self.backbone = network

        #spatial pooling layer
        self.pool = GeM(p = gem_p)

        #normalize on the unit-hypershpere
        self.norm = F.normalize

    def forward(self, img):
        '''
        Output has shape: batch size x descriptor dimension
        '''
        x = self.norm(self.pool(self.backbone(img))).squeeze(-1).squeeze(-1)

        return x

def extract_embeddings(net,dataloader,ms=[1],msp=1,print_freq=None,verbose = False):

    feature_extractor = create_feature_extractor(net, return_nodes=['backbone.getitem_5'])

    if verbose:

        if len(ms) == 1:
            print("Singlescale extraction")
        else:
            print("Multiscale extraction at scales: " + str(ms))
    
    net.eval()
    
    with torch.no_grad():

        vecs = np.zeros((768, len(dataloader.dataset))) #768 -&gt; vb16; 1024 -&gt; vl16
    
        for i,input in enumerate(dataloader):

            if len(ms) == 1 and ms[0] == 1:
                vecs[:, i] = feature_extractor(net,input[0])['backbone.getitem_5'] #.cuda()
            
            else:
                vecs[:, i] = feature_extractor(input[0])['backbone.getitem_5'] 

            if print_freq is not None:
                print("image: "+str(i))

    return vecs.T
</code></pre>
<p>But I get the error: TypeError: avg_pool2d(): argument ‘kernel_size’ (position 2) must be tuple of ints, not tuple. Does anyone know how I can fix it?</p>

username: ptrblck
cooked: <p>The error message sounds a bit strange but I guess the <code>dtype</code> of the kernel size might be a float. Could you check it and transform the values to <code>int</code>?</p>

username: whatevername
cooked: <p>If I try this: print(x.size(-2).dtype) it prints Proxy(getattr_16). Can I just put in (224,224) as the kernel size since those are the image dimensions?</p>

username: Matias_Vasquez
cooked: <p>I see that you have your function <code>extract_embeddings</code> but I cannot see where you use it.</p>
<p>In your <code>__init__</code> method inside your <code>Embedder</code> class you define your <code>backbone</code> as being this architecture, but then in your <code>forward</code> method you use the full model to process the image.</p>
<p>If this is really so, then at the end you should have a <code>1x1000</code> tensor with the predictions for each class.</p>
<p>Can you print the image of the output from the backbone like this ↓ ?</p>
<aside class="quote no-group quote-modified" data-username="whatevername" data-post="1" data-topic="151655">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/w/ebca7d/48.png" class="avatar"> whatevername:</div>
<blockquote>
<pre><code class="lang-auto">def forward(self, img):
        '''
        Output has shape: batch size x descriptor dimension
        '''
        x = self.backbone(img)
        print(x.shape) #
        x = self.norm(self.pool(x)).squeeze(-1).squeeze(-1)

        return x
</code></pre>
</blockquote>
</aside>

username: whatevername
cooked: <p>If I do that it prints this: Proxy(getattr_16).<br>
What I do is I initialize the model:</p>
<pre><code class="lang-auto">	if network_variant == 'vb16IN':
		net = Embedder("vit_b_16", pretrained_flag = True)
</code></pre>
<p>Then extract the embeddings like this:</p>
<pre><code class="lang-auto">train_descr = extract_embeddings(net,train_loader,ms = scales,msp = 1.0,print_freq=20000)
</code></pre>
<p>I tried to follow the feature extraction instruction here <a href="https://discuss.pytorch.org/t/assertionerror-expected-batch-size-seq-length-hidden-dim-got-torch-size-1-768-24-31/150869/4" class="inline-onebox">AssertionError: Expected (batch_size, seq_length, hidden_dim) got torch.Size([1, 768, 24, 31]) - #4 by Matias_Vasquez</a><br>
But I’m not sure I’m doing it right</p>

username: Matias_Vasquez
cooked: <h2><a name="shape-1" class="anchor" href="#shape-1"></a>Shape</h2>
<p>You get this ↓ by printing the shape of <code>x</code>?<br>
Can you print more information about <code>x</code>? Like the <code>type</code> or what it actually is?</p>
<aside class="quote no-group" data-username="whatevername" data-post="5" data-topic="151655">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/w/ebca7d/48.png" class="avatar"> whatevername:</div>
<blockquote>
<p>If I do that it prints this: Proxy(getattr_16).</p>
</blockquote>
</aside>
<pre data-code-wrap="python"><code class="lang-python">print(type(x))
print(x)
print(x.shape)
</code></pre>
<h2><a name="feature-extractor-2" class="anchor" href="#feature-extractor-2"></a>Feature Extractor</h2>
<p>What this <code>feature extractor</code> does, is run the input through the entire architecture and returns you a dictionary with the features that you request. In the example from that other response you can see that give the input image to the <code>feature extractor</code> and NOT to the <code>network</code>.</p>
<p>Maybe an easier way to put this by using another name. Here I only changed from <code>feature_extractor</code> to <code>network_using_features</code>.</p>
<pre data-code-wrap="python"><code class="lang-python"># This is the same code as I used in the other 
from torchvision.models.feature_extraction import create_feature_extractor

network = getattr(torchvision.models,"vit_b_16")(pretrained=True)
network_using_features = create_feature_extractor(network, return_nodes=['getitem_5'])

img = torch.rand(1, 3, 224, 224)
print(network_using_features(img)['getitem_5'])
</code></pre>
<hr>
<h2><a name="what-do-you-want-to-do-3" class="anchor" href="#what-do-you-want-to-do-3"></a>What do you want to do?</h2>
<p>But the main question is what you want to do with this architecture?<br>
I think it would really help if you explain what you are trying to do in simple words.</p>
<p>If you only want to fine-tune the Vision Transformer to classify images from your dataset, then we are overcomplicating stuff here.</p>
<p>Also, understanding how the architecture works and what it does to the image will really help you decide what/where/how you want to use the data coming from this architecture.</p>
<p>There are many great videos on youtube that explain this architecture.</p>

username: whatevername
cooked: <p>So what I want to do is get the output of getitem_5 and then apply gem pooling and normalization on it, is that possible? The bigger picture is: I want to extract the image descriptors using a Vision Transformer and then feed the descriptors into a kNN classifier.<br>
So how can I initialize the backbone to only be up until getitem_5? Because with ResNet I kept only the layers I wanted like this:</p>
<pre><code class="lang-auto">self.backbone = nn.Sequential(*list(network.children())[:-2])
</code></pre>
<p>But it’s not possible to do it like this with ViT</p>

username: whatevername
cooked: <p>I tried to fix it. Can you take a look at this code:</p>
<pre><code class="lang-auto">class GeM(nn.Module):
    
    def __init__(self, p=3, eps=1e-6):
        super(GeM,self).__init__()
        self.p = Parameter(torch.ones(1)*p)
        self.eps = eps

    def forward(self, x):
        return gem(x, p=self.p, eps=self.eps)
        
    def __repr__(self):
        return self.__class__.__name__ + '(' + 'p=' + '{:.4f}'.format(self.p.data.tolist()[0]) + ', ' + 'eps=' + str(self.eps) + ')'

def gem(x, p=3, eps=1e-6):
    print(x.size(-1))
    return F.avg_pool2d(x.clamp(min=eps).pow(p), (x.size(-2), x.size(-1))).pow(1./p)

class Embedder(nn.Module):

    def __init__(self,architecture,gem_p = 3,pretrained_flag = True):

        super(Embedder, self).__init__()

        #load the base model from PyTorch's pretrained models (imagenet pretrained)
        network = getattr(models,architecture)(pretrained=pretrained_flag)

        #keep only the convolutional layers, ends with relu to get non-negative descriptors
        if architecture.startswith('resnet'):
            self.backbone = nn.Sequential(*list(network.children())[:-2])
        elif architecture.startswith('efficientnet'):
            self.backbone = nn.Sequential(*list(network.children())[:-2])
        elif architecture.startswith('vit'):
            self.backbone = create_feature_extractor(network, return_nodes=['getitem_5'])

        #spatial pooling layer
        self.pool = GeM(p = gem_p)

        #normalize on the unit-hypershpere
        self.norm = F.normalize

        #information about the network
        self.meta = {
            'architecture' : architecture, 
            'pooling' : "gem",
            'mean' : [0.485, 0.456, 0.406], #imagenet statistics for imagenet pretrained models
            'std' : [0.229, 0.224, 0.225],
            'outputdim' : OUTPUT_DIM[architecture],
        }

    def forward(self, img):
        '''
        Output has shape: batch size x descriptor dimension
        '''
        x = self.norm(self.pool(self.backbone(img)['getitem_5'])).squeeze(-1).squeeze(-1)

        return x

def extract_ss(net, input):

    return net(input).cpu().data.squeeze()

def extract_ms(net, input, ms, msp):
    
    v = torch.zeros(net.meta['outputdim'])
    
    for s in ms:

        if s == 1:
            input_t = input.clone()
        
        else:    
            input_t = nn.functional.interpolate(input, scale_factor=s, mode='bilinear', align_corners=False)
        
        v += net(input_t).pow(msp).cpu().data.squeeze()
        
    v /= len(ms)
    v = v.pow(1./msp)
    v /= v.norm()

    return v

def extract_embeddings(net,dataloader,ms=[1],msp=1,print_freq=None,verbose = False):

    if verbose:

        if len(ms) == 1:
            print("Singlescale extraction")
        else:
            print("Multiscale extraction at scales: " + str(ms))
    
    net.eval()
    
    with torch.no_grad():

        vecs = np.zeros((net.meta['outputdim'], len(dataloader.dataset)))
    
        for i,input in enumerate(dataloader):

            if len(ms) == 1 and ms[0] == 1:
                vecs[:, i] = extract_ss(net,input[0]) #.cuda()
            
            else:
                vecs[:, i] = extract_ms(net,input[0], ms, msp) #.cuda()

            if print_freq is not None:
                if i%print_freq == 0:
                     print("image: "+str(i))

    return vecs.T
</code></pre>
<p>But now I get the error:</p>
<blockquote>
<p>IndexError: Dimension out of range (expected to be in range of [-2, 1], but got -3)</p>
</blockquote>
<p>from line:</p>
<pre><code class="lang-auto">return F.avg_pool2d(x.clamp(min=eps).pow(p), (x.size(-2), x.size(-1))).pow(1./p)
</code></pre>

username: Matias_Vasquez
cooked: <h2><a name="new-approach-1" class="anchor" href="#new-approach-1"></a>New approach</h2>
<p>So what you did now looks way better. Now you are really using the feature extractor to get some information from the architecture and use it. (and not the whole architecture).</p>
<p>How you may have seen, there are still some problems.</p>
<p>In order to fix them we need to truly understand what is happening with ViT and what we want to do. (That is why I asked what you intended to do with this.)</p>
<p>If your approach is to simply use ViT and classify directly with it, then there is a simpler approach. If as you mentioned, you want to use the “image features” and feed them to another architecture, then we need to understand ViT.</p>
<ul>
<li>ViT as classifier for N classes</li>
<li>Use “image features” and feed them to something else</li>
</ul>
<h2><a name="vit-as-classifier-for-n-classes-2" class="anchor" href="#vit-as-classifier-for-n-classes-2"></a>ViT as classifier for N classes</h2>
<p>Here there is not much to understand. We have our architecture and in the end there is a head that classifies the image into <code>N</code> classes. To change this we only need to change the head to meet our needs and fine-tune the model.</p>
<p>For the <code>torchvision</code> implementation of <code>vit_b_16</code>, we can access the head like this</p>
<pre data-code-wrap="python"><code class="lang-python">import torchvision

model = torchvision.models.vit_b_16(pretrained=True)
print(model.heads)

# Output:
#Sequential(
#  (head): Linear(in_features=768, out_features=1000, bias=True)
#)
</code></pre>
<p>So if we wanted to change this it would be very straighforward.</p>
<p>With this simple code we get our model with the right amount of classes and we can fine-tune it to classify our specific dataset.</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torchvision

N_CLASSES = 2

model = torchvision.models.vit_b_16(pretrained=True)
model.heads.head = torch.nn.Linear(768, N_CLASSES)
</code></pre>
<h2><a name="use-image-features-and-feed-them-to-something-else-3" class="anchor" href="#use-image-features-and-feed-them-to-something-else-3"></a>Use “image features” and feed them to something else</h2>
<p>If we want to do this approach, then we need to first understand ViT.</p>
<h3><a name="understanding-vit-4" class="anchor" href="#understanding-vit-4"></a>Understanding ViT</h3>
<p>So these are the steps that we will look into:</p>
<ul>
<li>Making patches</li>
<li>Flatten the patches</li>
<li>Class token</li>
<li>Positional Embedding</li>
<li>Encoder
<ul>
<li>Block</li>
</ul>
</li>
<li>Head</li>
</ul>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/f/ff3498558c13bf384b5fed20fae569d663b200d9.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/ff3498558c13bf384b5fed20fae569d663b200d9" title="grafik"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/f/ff3498558c13bf384b5fed20fae569d663b200d9_2_690x350.jpeg" alt="grafik" data-base62-sha1="ApEreRCpZ9MuPuiKdVE2m0LHguZ" width="690" height="350" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/f/ff3498558c13bf384b5fed20fae569d663b200d9_2_690x350.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/f/f/ff3498558c13bf384b5fed20fae569d663b200d9_2_1035x525.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/f/f/ff3498558c13bf384b5fed20fae569d663b200d9_2_1380x700.jpeg 2x" data-dominant-color="E3E2E0"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">grafik</span><span class="informations">1400×712 151 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<ul>
<li>Making patches</li>
</ul>
<p>On the lower left corner of the image you can see how the image is divided into patches. For the architecture we are using (<code>vit_b_16</code>) each patch has <code>3x16x16</code> pixels. For an image with <code>224</code> pixels, we get <code>14x14</code> patches (<code>14*16=224</code>). This is why the input image has to have this size.</p>
<p>If we have a batch (<code>B</code>) of one image we could write it like this <code>B x Patch_j x Patch_i x C x H x W = 1 x 14 x 14 x 3 x 16 x 16</code>.</p>
<p>This means we have one image. This image has 14 patches in the y direction (j) and 14 patches in the x direction (i). Each of these patches have 3 channels. These patches also have 16 by 16 pixels in the x and y directions.</p>
<ul>
<li>Flatten the patches</li>
</ul>
<p>In the image above, we can see there is a pink box with the title <code>Linear Projection of Flattened Patches</code>. What we are making is rearanging these patches. So before we had this structure</p>
<p><code>B x Patch_j x Patch_i x C x H x W = 1 x 14 x 14 x 3 x 16 x 16</code></p>
<p>now  we merge the <code>14x14</code> patches into one dimension. We also merge <code>CxHxW</code> into one direction.</p>
<p><code>B x Patches x Pixels = 1 x 14*14 x 3*16*16 = 1 x 196 x 768</code></p>
<p>This is also known as <code>Batch x Seq_Length x Hidden_Dim</code>.</p>
<p>So if we understand this, we can see that here the image has another meaning.</p>
<ul>
<li>Class token</li>
</ul>
<p>After this pink box on the architecture, there are several pink boxes with numbers and an extra box with an <em>asterix</em>  on the left. As described on the image, this asterix is the Extra learnable [class] embedding. Here is where the results of the image transformer will be.</p>
<p>The size of this class token is <code>B x 1 x Hidden_Dim = 1 x 1 x 768</code> for our case.</p>
<p>This is then prepended at the beggining of our representation of the image on the <code>Seq_Length</code> dim.</p>
<p><code>[cls]; img -&gt; B x Seq_length + 1 x Hidden_Dim = 1 x 197 x 768</code>.</p>
<p>Now <code>Seq_length = 197</code>. Here we have our <code>14x14</code> patches plus one for our [cls] token.</p>
<ul>
<li>Positional Embedding</li>
</ul>
<p>Also on the image we can see that a positional embedding is added. This is in order to see how the position of the patches relate to eachother.</p>
<p>The Dimentions however, stay the same.</p>
<p><code>B x Seq_Length x Hidden_Dim = 1 x 197 x 768</code></p>
<ul>
<li>Encoder</li>
</ul>
<p>Now come the fun part. On the right of the image is the architecture of ONE encoder block. There are <code>L</code> blocks in the ViT architecture. (<code>L=12</code> in ours). Meaning the image is fed to the first encoder block and the output of the first comes to the second and third sequentially.</p>
<p>All of the blocks have the same architecture, meaning the size of the output has to be the same as input in order to be fed to the next block.</p>
<p><code>Size_in  = B x Seq_Length x Hidden_Dim = 1 x 197 x 768</code><br>
<code>Size_out = B x Seq_Length x Hidden_Dim = 1 x 197 x 768</code></p>
<ul>
<li>Block</li>
</ul>
<p>If we look at how the block architecture is built, we can see some skip layers, batch normalizations and a multi-layer perceptron (MLP) in the end. There is also a Multi-Head Attention Layer (MHA).</p>
<p>The MHA consists of multiple parallel scaled dot product attention mechanisms with learnable parameters.</p>
<p>One Scaled dot product attention looks like this<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/5/15095597e0a697754bd8c90c5427f357e16bfbad.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/15095597e0a697754bd8c90c5427f357e16bfbad" title="grafik"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/5/15095597e0a697754bd8c90c5427f357e16bfbad_2_233x250.jpeg" alt="grafik" data-base62-sha1="3061Awbv2ea5Ld2sztvCkpDI62F" width="233" height="250" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/5/15095597e0a697754bd8c90c5427f357e16bfbad_2_233x250.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/1/5/15095597e0a697754bd8c90c5427f357e16bfbad_2_349x375.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/1/5/15095597e0a697754bd8c90c5427f357e16bfbad_2_466x500.jpeg 2x" data-dominant-color="E9E8E8"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">grafik</span><span class="informations">684×732 36.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Described by the following equation</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/0/7/07adc4f1f3784b969a200fda0196c210a7fd091f.png" alt="grafik" data-base62-sha1="15VDLiGTuBgsbdldurnyn7tI3YX" width="450" height="60"></p>
<p>There is too much to unpack about this equation, but the important thing is that the pixels are attending to themselves. And <em>HERE</em> might be a good place to get the image features from.</p>
<p>To see how, see below in the <code>Image Features</code> section.</p>
<ul>
<li>Head</li>
</ul>
<p>Up until now our data has the following format</p>
<p><code>B x CLS_Token+Seq_Length x Hidden_dim = 1 x 197 x 768</code>.</p>
<p>We said that the CLS Token is where the classification will be done. So now we only take the CLS token and do not care for the rest.</p>
<p><code>B x CLS_Token x Seq_Length = 1 x 1 x 768</code></p>
<p>As we saw on the beggining of this post, the head consists (in this case) of a <code>Linear</code> layer that has <code>768</code> in features and <code>1000</code> features. This means that we feed our data to this <code>Linear</code> layer and we classify between 1000 classes.</p>
<p>On the image with the ViT architecture, this is represented by the yellow box with <code>MLP Head</code> written on it.</p>
<h3><a name="image-features-5" class="anchor" href="#image-features-5"></a>Image Features</h3>
<p><a href="https://github.com/sayakpaul/probing-vits" rel="noopener nofollow ugc">Here</a> is an implementation of tensorflow to get the attention map.</p>
<p><a href="https://www.youtube.com/watch?v=I3VnY01JpSo" rel="noopener nofollow ugc">Here</a> is a video + code on how to get them for pytorch.</p>
<p><a href="https://arxiv.org/pdf/2108.08810.pdf" rel="noopener nofollow ugc">Here</a> is also a paper that might interest you.</p>
<p>These methods however, require that you either switch to TF or rewrite the ViT to access the attention map.</p>
<p>With feature extractor you can get intermediate steps (as we have already done by taking the <code>getitem_5</code>, which is almost the last step in the architecture).</p>
<p>If we inspect the <code>graph_node_names</code> as we did before to get the <code>getitem_5</code> name, we can see that there is a <code>encoder.layers.encoder_layer_11.self_attention</code>. But the result is after doing the full Multi-Head Self-Attention.</p>
<p>We want an intermediate result.<br>
So for this approach you could do something like this. (I ran this in a python notebook to see all of the heads.)</p>
<pre data-code-wrap="python"><code class="lang-python">import matplotlib.pyplot as plt
import seaborn as sns
from PIL import Image
import requests
from io import BytesIO
import torchvision.transforms as T
from torchvision.models.feature_extraction import create_feature_extractor, get_graph_node_names


url = "https://www.thesprucepets.com/thmb/k3NXIqobAKvxoQ2ozGcwPxzIkpI=/3300x1856/smart/filters:no_upscale()/most-obedient-dog-breeds-4796922-hero-4440a0ccec0e42c98c5e58821fc9f165.jpg"
response = requests.get(url)
img = Image.open(BytesIO(response.content))

img = T.Resize((224,224))(T.ToTensor()(img))
plt.imshow(img.permute(1, 2, 0))
plt.show()


model = torchvision.models.vit_b_16(pretrained=True)
keys = ['encoder.layers.encoder_layer_11.ln']
feature_extractor = create_feature_extractor(model, return_nodes=keys)
feature_extractor.eval()


out = feature_extractor(img.unsqueeze(0))

x = out[keys[0]]
x_, attn = model.encoder.layers[-1].self_attention(x, x, x, need_weights=True, average_attn_weights=False)

print(x_.shape)
print(attn.shape)

for i in range(12):
    sns.heatmap(attn[0, i, 0, 1:].view(14, 14).detach().numpy())
    plt.show()
</code></pre>
<p>The output will be the <code>12</code> self attention heads ploted as heatmaps. As you will see, for this particular example, many will mean nothing to you but one of the heads looks like this.<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/1/2/12b787eeed126891102301f3aed870303a4e2dea.png" alt="grafik" data-base62-sha1="2FzO5bfwMxAHC2FIj3qVCsBjtWO" width="192" height="189"><br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/5/8/580515dfbdd7a4c6687758dbe9ab689eb46681b3.png" alt="grafik" data-base62-sha1="cyEUBfo9ucwKatOE2tmOYhMC08j" width="271" height="186"></p>
<p>If you set <code>average_attn_weigths=True</code> you will get the average of all 12 attention heads and will also mean nothing to us.</p>
<p>But if you also get the ‘heads’ key and print the predicted class, you will see that the predicted class is correct (208=Labrador retriever). So it means that it is working.<br>
<code>keys = ['encoder.layers.encoder_layer_11.ln', 'heads']</code><br>
<code>print(out['heads'].argmax(dim=1))</code></p>
<p>You could do something like this and use this self attention given by the ViT, but you need to understand what is happening and what you want to use and how.</p>
<p>Also, this will only be a suggestion of where important features might be. But these alone may mean nothing when feeding them to another architecture. So you might want to expriment by scaling them to the actual size of the image and feed both the image AND this heatmap to another architecture.</p>
<p>But these are just suggestions of what you could theretically do.</p>
<p>The most important thing is understanding what it does and how it does it. Then you can decide how to proceede with this information.</p>
<p>Hope this is a bit clearer now.</p>

username: Matias_Vasquez
cooked: <p>Oh and I forgot to address the error that you mentioned.</p>
<aside class="quote no-group quote-modified" data-username="whatevername" data-post="8" data-topic="151655">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/w/ebca7d/48.png" class="avatar"> whatevername:</div>
<blockquote>
<p>But now I get the error:</p>
<blockquote>
<p>IndexError: Dimension out of range (expected to be in range of [-2, 1], but got -3)</p>
</blockquote>
<p>from line:</p>
<pre><code class="lang-auto">return F.avg_pool2d(x.clamp(min=eps).pow(p), (x.size(-2), x.size(-1))).pow(1./p)
</code></pre>
</blockquote>
</aside>
<p>As you now know, the <code>getitem_5</code> stage of the ViT is right before the <code>head</code>. Right after we selected only the CLS Token. Which we mentioned here ↓</p>
<aside class="quote no-group" data-username="Matias_Vasquez" data-post="9" data-topic="151655">
<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>We said that the CLS Token is where the classification will be done. So now we only take the CLS token and do not care for the rest.</p>
<p><code>B x CLS_Token x Seq_Length = 1 x 1 x 768</code></p>
</blockquote>
</aside>
<p>If you see, the shape here is <code>Bx1x768</code>. If you try to do <code>avg_pool2d</code> with this tensor, which has this weird shape then it will give an error.</p>
<p>The expected input for <code>avg_pool2d</code> has the shape <code>B x C x H x W</code>.</p>

username: whatevername
cooked: <p>Thank you so much for the very clear explanation! I’ll try to implement all of this.</p>