pytorch利用resnet50提取图像特征

我不是女神ヾ 2022-10-28 00:58 258阅读 0赞

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()语句，会导致显存溢出从而报错退出。