pytorch使用预训练图像模型分类预测

指定设备

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

工具包

import os

import cv2
from PIL import Image, ImageFont, ImageDraw

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
%matplotlib inline

import torch
from torchvision import models
import torch.nn.functional as F
from torchvision import transforms

加载预训练模型

从torchvision导入的models中包含了许多预训练模型

model = models.resnet18(pretrained=True) 

model = model.eval()
model = model.to(device)

图像预处理

transforms从torvhvison导入
RCTN模板：缩放、裁剪、转tensor、归一化（均值标准差按照惯例）

test_transform = transforms.Compose([transforms.Resize(256),
                                     transforms.CenterCrop(224),
                                     transforms.ToTensor(),
                                     transforms.Normalize(
                                         mean=[0.485, 0.456, 0.406], 
                                         std=[0.229, 0.224, 0.225])
                                    ])

导入图像

使用pillow

img_pil = Image.open(img_path) # 用 pillow 载入

执行预测

unsqueeze增加一个维度后才可传入模型。
可以看到这里的test_transform是一个对象，也就是经过上述compose后的返回值。

input_img = test_transform(img_pil).unsqueeze(0).to(device) # 预处理

在模型输出得到类别的分数，由于需要置信度，需要使用softmax归一化，结果是所有类别的和为1，并且得到了置信度。

pred_softmax = F.softmax(pred_logits, dim=1) # 对 logit 分数做 softmax 运算

可视化

关键是y传入时的调用

plt.figure(figsize=(8,4))

x = range(1000)
y = pred_softmax.cpu().detach().numpy()[0]

ax = plt.bar(x, y, alpha=0.5, width=0.3, color='yellow', edgecolor='red', lw=3)
plt.ylim([0, 1.0]) # y轴取值范围

获取top k结果，注意类别与置信度的解析

n = 10的
top_n = torch.topk(pred_softmax, n) # 取置信度最大的 n 个结果
pred_ids = top_n[1].cpu().detach().numpy().squeeze() # 解析出类别
confs = top_n[0].cpu().detach().numpy().squeeze() # 解析出置信度