pytorch利用resnet50提取图像特征

import os
import torch
import torch.nn as nn
from torchvision import models,transforms
from torch.autograd import Variable
import numpy as np
from PIL import Image
import time

device=torch.device("cuda:3" if torch.cuda.is_available() else "cpu")

class net(nn.Module):
    def __init__(self):
        super(net, self).__init__()
        self.net = models.resnet50(pretrained=True)

    def forward(self, input):
        output = self.net.conv1(input)
        output = self.net.bn1(output)
        output = self.net.relu(output)
        output = self.net.maxpool(output)
        output = self.net.layer1(output)
        output = self.net.layer2(output)
        output = self.net.layer3(output)
        output = self.net.layer4(output)
        output = self.net.avgpool(output)
        return output

features_dir="./index/"
imgPath="./test/2cf80bb7a9fb5c63c20401e8722408f0.jpg"
fileName=imgPath.split('/')[-1]
feature_path=os.path.join(features_dir,fileName+".txt")

transform = transforms.Compose([
    # transforms.Resize(256),
    # transforms.CenterCrop(224),
    transforms.ToTensor()]
)

img = Image.open(imgPath)
img = transform(img)
print(img.shape)
startTime=time.perf_counter()
x = Variable(torch.unsqueeze(img, dim=0).float(), requires_grad=False)
print(x.shape)
with torch.no_grad():
    model=net()
    # if use_gpu:
    #     x = x.cuda()
    #     net = net.cuda()
    y = model(x).cpu()
    y = torch.squeeze(y)
    y = y.data.numpy()
    print(y.shape)
endTime=time.perf_counter()
print("cpu time:"+str(endTime-startTime))
model2=model.cuda()
endTime=time.perf_counter()
with torch.no_grad():
    x = Variable(torch.unsqueeze(img, dim=0).float(), requires_grad=False)
    x=x.cuda()
    y = model2(x).cpu()
    y = torch.squeeze(y)
    y = y.data.numpy()
    print(y.shape)
endTime2=time.perf_counter()
print("GPU time:"+str(endTime2-endTime))
np.savetxt(feature_path, y, delimiter=',')

需要注意：如果没有使用 with torch.no_grad()语句，会导致显存溢出从而报错退出。