Hy guys, I want to extract the in_feature(2048) of FC layer, passing an image to resnet50.
def get_vector(image):
layer = model._modules.get('avgpool')
my_embedding = torch.zeros(2048) #2048 is the in_features of FC , output of avgpool
def copy_data(m, i, o):
my_embedding.copy_(o.data)
h = layer.register_forward_hook(copy_data)
model(image)#error
h.remove()
# return the vector
return my_embedding
from PIL import Image
from torch.utils.data import DataLoader
from torchvision import transforms
from os import listdir
from os import path
import time
import h5py
import pandas as pd
since = time.time()
print("---- START FEATURES EXTRACTION ----")
print
column = ["FlickrID", "Features"]
path = "./train_dataset/train_imgs/"
pathCSV = "./train_dataset/features/img_info_TRAIN.csv"
f_id=[]
features_extr=[]
df = pd.DataFrame(columns=column)
tr=transforms.Compose([transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
test = Dataset(path, pathCSV, transform=tr)
test_loader = DataLoader(test, batch_size=1, num_workers=2, shuffle = False)
for batch in test_loader:
nome = batch['FlickrID']
f_id.append(nome)
image = batch['image']
with torch.no_grad():
pred = get_vector(image)#error
features_extr.append(pred)
df["FlickrID"] = f_id
df["Features"] = features_extr
df.to_hdf("Places.h5", key='df', mode='w')
time_elapsed = time.time() - since
print('---- FEATURES SAVED in {:.0f}m {:.0f}s ----'.format(time_elapsed // 60, time_elapsed % 60))
I have this error: output with shape [2048] doesn’t match the broadcast shape [1, 2048, 1, 2048] in the line of code that I have commented with #error
Can you help me?