从 pytorch 中预训练的 resnet50 中提取特征

Question

大家好，我想提取预训练 resnet50 的 in_features 全连接层。

我之前创建了一个给我特征向量的方法：

def get_vector(image):

#layer = model._modules.get('fc')

layer = model.fc
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)
tmp = model(image)

h.remove()

# return the vector
return my_embedding

我在这里调用这个方法后：

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)



#Leggiamo le immagini
for batch in test_loader:
    nome = batch['FlickrID']
    f_id.append(nome)
    image = batch['image']



    #print(image)
    with torch.no_grad():
        pred = get_vector(image)

    features_extr.append(pred)

df["FlickrID"] = f_id
df["Features"] = features_extr  


df.to_hdf("Places.h5", key='df', mode='w')

我遇到这样的错误：形状 [2048] 的输出与广播形状 [1, 2048, 1, 2048]

不匹配

这个resnet50的Fully Connected怎么取in_feature？数据集是自定义数据集 class.

抱歉我的英语不好

Answer 1

该模型采用批量输入，这意味着全连接层的输入大小为 [batch_size、2048]。因为您使用的批量大小为 1，所以它变为 [1, 2048]。因此，它不适合张量 torch.zeros(2048)，所以它应该是 torch.zeros(1, 2048)。

您还尝试使用层 model.fc 的输出 (o) 而不是输入 (i)。

除此之外，使用钩子过于复杂，获得特征的更简单方法是通过将 model.fc 替换为 nn.Identity 来修改模型，这只是 returns输入作为输出，由于特征是它的输入，所以整个模型的输出将是特征。

model.fc = nn.Identity()

features = model(image)

从 pytorch 中预训练的 resnet50 中提取特征

Extract features from pretrained resnet50 in pytorch

python

pytorch

resnet