Блог

У меня есть следующий код, и я пытаюсь извлечь вектор измерения 2048 из предварительно обученной сети resnet50.

 import torch
import torchvision
import torchvision.models as models
from PIL import Image
# Load the pretrained model
model = models.resnet50(pretrained=True)

# Use the model object to select the desired layer
layer = model._modules.get('avgpool')

# Set model to evaluation mode
model.eval()

transforms = torchvision.transforms.Compose([
    torchvision.transforms.Resize(256),
    torchvision.transforms.CenterCrop(224),
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])


def get_vector(image):
    # Create a PyTorch tensor with the transformed image
    t_img = transforms(image)
    
    # 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)
    my_embedding = torch.zeros(2048)


    # Define a function that will copy the output of a layer
    def copy_data(m, i, o):
        my_embedding.copy_(o.flatten())                 # <-- flatten

    # Attach that function to our selected layer
    h = layer.register_forward_hook(copy_data)
    # Run the model on our transformed image
    with torch.no_grad():                               # <-- no_grad context
        model(t_img.unsqueeze(0))                       # <-- unsqueeze
    # Detach our copy function from the layer
    h.remove()
    # Return the feature vector
    return my_embedding
import torch.nn as nn
cos = nn.CosineSimilarity(dim=1, eps=1e-6)
img1 = Image.open('1.jpg')
img2 = Image.open('2.jpg')
img3 = Image.open('3.jpg')
output1 = cos(pic_vector1, pic_vector2)
  

Я получаю следующую ошибку:

 ---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-19-88106798854f> in <module>
----> 1 output1 = cos(pic_vector1, pic_vector2)

~/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    720             result = self._slow_forward(*input, **kwargs)
    721         else:
--> 722             result = self.forward(*input, **kwargs)
    723         for hook in itertools.chain(
    724                 _global_forward_hooks.values(),

~/anaconda3/lib/python3.7/site-packages/torch/nn/modules/distance.py in forward(self, x1, x2)
     73 
     74     def forward(self, x1: Tensor, x2: Tensor) -> Tensor:
---> 75         return F.cosine_similarity(x1, x2, self.dim, self.eps)

IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)
  

Как я могу это исправить?