用pytorch进行图像分类预测

提示：以下是本篇文章正文内容，本文是通过阅读TommyZihao在Githup的资源所写，有什么不对的地方欢迎指正。

目录

前言

一、环境的安装

二、应用模型对图片进行分类

---

前言

本篇文章的主要内容是，通过pytorch中现有的数据分类模型，调用pytorch中现有的模型进行图片，视频，摄像头实时画面的预测分类。

---

提示：以下是本篇文章正文内容，仅供参考

一、环境的安装

1.下载pytorch

!pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113

 上述在代码在jupyter Notebook中可直接运行，如果运行上述程序中出现报错可以考虑是‘’网络问题‘’，建议更换更稳定的网络试试，若还未解决那就百度吧。

2.下载中文字体

 下载地址：https://us-logger1.oss-cn-beijing.aliyuncs.com/SimHei.ttf

 3.数据集

 数据集官网下载链接：https://image-net.org/

二、应用模型对图片进行分类

# 有 GPU 就用 GPU，没有就用 CPU
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

 运行后若显示‘CPU’，则为没有启用GPU只启用了CPU。

from torchvision import model

# 载入预训练图像分类模型
model = models.resnet18(pretrained=True) 

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

首先把我们所需要的训练分类的模型导入，现在所采用的分类模型为resnet18这个模型深度为18，如想获得更高的精度可改用152.

from torchvision import transforms

# 测试集图像预处理-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])
                                    ])

对图像进行预处理，这是非常重要的一步，需要对图片进行缩放，调整等。

img_path = 'test_img/dog.jpeg'
# 用 pillow 载入
from PIL import Image
img_pil = Image.open(img_path)
#转换成array形式的数组
np.array(img_pil).shape
# 预处理
input_img = test_transform(img_pil) 
input_img = input_img.unsqueeze(0).to(device)
# 执行前向预测，得到所有类别的 logit 预测分数
pred_logits = model(input_img) 
import torch.nn.functional as F
 # 对 logit 分数做 softmax 运算
pred_softmax = F.softmax(pred_logits, dim=1) 

 通过上述程序，获得了每个别的置信的，置信度越大那么证明该图片预测属于哪一类。

n = 4
top_n = torch.topk(pred_softmax, n)
#运行结果
#torch.return_types.topk(
#values=tensor([[0.7844, 0.1232, 0.0022, 0.0009]], device='cuda:0', grad_fn=<TopkBackward0>),
#indices=tensor([[233, 234, 238, 390]], device='cuda:0'))

# 解析出类别
pred_ids = top_n[1].cpu().detach().numpy().squeeze()
#运行结果
#array([233, 234, 238, 390])

# 解析出置信度
confs = top_n[0].cpu().detach().numpy().squeeze()
#运行结果
#array([0.78444356, 0.12316445, 0.00216345, 0.00094444],dtype=float32)

 

---