使用pytorch預訓練模型分類與特征提取

    pytorch(pytorch v0.1 這個是早期版本了)應該是深度學習架構裡面比較好使用的了，相比于tensorflow，mxnet。可能在使用者上稍微少一點，有的時候出問題不好找文章。下面就使用pytorch預訓練模型做分類和特征提取，pytorch文檔可以參考：pytorch docs  , 模型是imagenet2012訓練的标簽可參考：imagenet2012 labels  ，模型預測的下标按從上到下，起始（n01440764）為0

#encoding=utf-8

import os
import numpy as np

import torch
import torch.nn
import torchvision.models as models
from torch.autograd import Variable 
import torch.cuda
import torchvision.transforms as transforms

from PIL import Image

transform_list = [transforms.ToTensor(),
                  transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                       std =[0.229, 0.224, 0.225])]
img_to_tensor = transforms.Compose(transform_list)


def make_model():
    resmodel=models.resnet34(pretrained=True)
    resmodel.cuda()#将模型從CPU發送到GPU,如果沒有GPU則删除該行
    return resmodel

#分類
def inference(resmodel,imgpath):
    resmodel.eval()#必需，否則預測結果是錯誤的
    
    img=Image.open(imgpath)
    img=img.resize((224,224))
    tensor=img_to_tensor(img)
    
    tensor=tensor.resize_(1,3,224,224)
    tensor=tensor.cuda()#将資料發送到GPU，資料和模型在同一個裝置上運作
            
    result=resmodel(Variable(tensor))
    result_npy=result.data.cpu().numpy()#将結果傳到CPU，并轉換為numpy格式
    max_index=np.argmax(result_npy[0])
    
    return max_index
    
#特征提取
def extract_feature(resmodel,imgpath):
    resmodel.fc=torch.nn.LeakyReLU(0.1)
    resmodel.eval()
    
    img=Image.open(imgpath)
    img=img.resize((224,224))
    tensor=img_to_tensor(img)
    
    tensor=tensor.resize_(1,3,224,224)
    tensor=tensor.cuda()
            
    result=resmodel(Variable(tensor))
    result_npy=result.data.cpu().numpy()
    
    return result_npy[0]
    
if __name__=="__main__":
    model=make_model()
    imgpath='path_to_img/xxx.jpg'
    print inference(model,imgpath)
    print extract_feature(model, imgpath)
    
           

參考：  pytorch doc