说明
最近在研究分类模型如何部署C++，先拿Resnet50 来练一练手，文章将 分为多篇，这一篇主要验证一下pytorch 模型输出是正确的，为后续tensort RT 模型输出提供验证。
1 官方权重下载
https://download.pytorch.org/models/resnet50-19c8e357.pth
2 测试代码
import torchvision.models as models
import os
import json
import torch
from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as pltdef main():# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")device = torch.device("cuda:0")# device = torch.device("cpu")model = models.resnet50()model = model.cuda()# num_classes = 10 # 修改为你自己的类别数量# model.fc = torch.nn.Linear(model.fc.in_features, num_classes)data_transform = transforms.Compose([transforms.Resize(256),transforms.CenterCrop(224),transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])# model.load_state_dict(torch.load('params.pth', map_location=device)) #model.load_state_dict(torch.load('resnet50-19c8e357.pth'))model.eval()img_path = 'data/fox.png'assert os.path.exists(img_path), "file: '{}' dose not exist.".format(img_path)img = Image.open(img_path)plt.imshow(img)# [N, C, H, W]img = data_transform(img)# expand batch dimensionimg = torch.unsqueeze(img, dim=0)with torch.no_grad():# predict classoutput = torch.squeeze(model(img.to(device))).cpu()predict = torch.softmax(output, dim=0)print(predict.shape)print(predict[270:280]) # 打印几个softmax之后的输出predict_cla = torch.argmax(predict).numpy() # 找出最大的序号print(predict_cla) # 打印出类别if __name__ == '__main__':main()
3 测试图片
这是我找的几张图片
4 测试结果
跑一张狐狸的图片，输出序号277
查看预训练模型对应的类别编号
