RuntimeError：使用 pytorch 提取特征向量时，形状为 [512] 的输出与广播形状 [1, 512, 1, 512] 不匹配

Question

我无法解决此错误。此代码取自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)

错误：-

---------------------------------------------------------------------------
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]

我实际上想要做的是尝试从我想进一步用于构建推荐系统的图像中提取特征向量。如果有其他可用的替代方案，请通知我。 先谢谢了！！！

Answer 1

或者，您可以将钩子函数替换为以下内容，这样您就不必处理调整输出形状的问题：

# step 3 and 4
my_embedding = None
def my_hook(module_, input_, output_):
    nonlocal my_output
    my_embedding = output_

然后只需调用以下代码

# step 5
h = layer.register_forward_hook(my_hook)

Answer 2

avgpool之后需要reshape输出数据：

def copy_data(m, i, o):
    my_embedding.copy_(o.data.reshape(o.data.size(1)))