RuntimeError: output with shape [512] doesn't match the broadcast shape [1, 512, 1, 512] while extracting feature vector using pytorch - machine-learning

Question

I am not able to resolve this error. This code is taken from https://becominghuman.ai/extract-a-feature-vector-for-any-image-with-pytorch-9717561d1d4c
import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from torch.autograd import Variable
from PIL import Image
pic_one = '/content/drive/My Drive/Video_Recommender/zframe1.jpg'
pic_two = '/content/drive/My Drive/Video_Recommender/zframe2.jpg'
model = models.resnet18(pretrained=True)
layer = model._modules.get('avgpool')
scaler = transforms.Scale((224, 224))
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
to_tensor = transforms.ToTensor()
def get_vector(image_name):
# 1. Load the image with Pillow library
img = Image.open(image_name)
# 2. Create a PyTorch Variable with the transformed image
t_img = Variable(normalize(to_tensor(scaler(img))).unsqueeze(0))
# 3. Create a vector of zeros that will hold our feature vector
# The 'avgpool' layer has an output size of 512
my_embedding = torch.zeros(512)
# 4. Define a function that will copy the output of a layer
def copy_data(m, i, o):
my_embedding.copy_(o.data)
# 5. Attach that function to our selected layer
h = layer.register_forward_hook(copy_data)
# 6. Run the model on our transformed image
model(t_img)
# 7. Detach our copy function from the layer
h.remove()
# 8. Return the feature vector
return my_embedding
pic_one_vector = get_vector(pic_one)
pic_two_vector = get_vector(pic_two)
Error:-
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
<ipython-input-41-ca2d66de2d9c> in <module>()
----> 1 pic_one_vector = get_vector(pic_one)
2 pic_two_vector = get_vector(pic_two)
5 frames
<ipython-input-40-a45affe9d8f7> in get_vector(image_name)
13 h = layer.register_forward_hook(copy_data)
14 # 6. Run the model on our transformed image
---> 15 model(t_img)
16 # 7. Detach our copy function from the layer
17 h.remove()
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
548 result = self._slow_forward(*input, **kwargs)
549 else:
--> 550 result = self.forward(*input, **kwargs)
551 for hook in self._forward_hooks.values():
552 hook_result = hook(self, input, result)
/usr/local/lib/python3.6/dist-packages/torchvision/models/resnet.py in forward(self, x)
218
219 def forward(self, x):
--> 220 return self._forward_impl(x)
221
222
/usr/local/lib/python3.6/dist-packages/torchvision/models/resnet.py in _forward_impl(self, x)
211 x = self.layer4(x)
212
--> 213 x = self.avgpool(x)
214 x = torch.flatten(x, 1)
215 x = self.fc(x)
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
550 result = self.forward(*input, **kwargs)
551 for hook in self._forward_hooks.values():
--> 552 hook_result = hook(self, input, result)
553 if hook_result is not None:
554 result = hook_result
<ipython-input-40-a45affe9d8f7> in copy_data(m, i, o)
9 # 4. Define a function that will copy the output of a layer
10 def copy_data(m, i, o):
---> 11 my_embedding.copy_(o.data)
12 # 5. Attach that function to our selected layer
13 h = layer.register_forward_hook(copy_data)
RuntimeError: output with shape [512] doesn't match the broadcast shape [1, 512, 1, 512]
What I am actually trying to do is trying to extract feature vector from images that I want to further use for building a recommendation system. Do inform me if there is any other alternative available.
Thanks in advance!!!

Answer 1

You need to reshape the output data after avgpool:
def copy_data(m, i, o):
my_embedding.copy_(o.data.reshape(o.data.size(1)))

Answer 2

Alternatively, you may replace the hook function with the following, just so you don't have to deal with adjusting the output shape:
# step 3 and 4
my_embedding = None
def my_hook(module_, input_, output_):
nonlocal my_output
my_embedding = output_
then simply call the following
# step 5
h = layer.register_forward_hook(my_hook)