Resnet C ++ 部署 pytorch功能测试(一)

说明

最近在研究分类模型如何部署C++,先拿Resnet50 来练一练手,文章将 分为多篇,这一篇主要验证一下pytorch 模型输出是正确的,为后续tensort RT 模型输出提供验证。

1 官方权重下载

https://download.pytorch.org/models/resnet50-19c8e357.pth

2 测试代码

python 复制代码
import torchvision.models as models
import os
import json
import torch
from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt


def 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 dimension
    img = torch.unsqueeze(img, dim=0)
    with torch.no_grad():
        # predict class
        output = 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

查看预训练模型对应的类别编号

相关推荐
wkj0011 分钟前
php中array($this, ‘loadClass‘)表示啥意思?
android·开发语言·php
Stuomasi_xiaoxin2 分钟前
服务器重装后如何“复活”旧硬盘上的 Anaconda 环境?—— 一次完整的排错与恢复记录
开发语言·python·github
presenttttt44 分钟前
用Python和OpenCV从零搭建一个完整的双目视觉系统(四)
开发语言·python·opencv·计算机视觉
爱学习的小邓同学1 小时前
数据结构 --- 队列
c语言·数据结构
每日出拳老爷子1 小时前
[C#] 使用TextBox换行失败的原因与解决方案:换用RichTextBox的实战经验
开发语言·c#
半桔1 小时前
【Linux手册】从接口到管理:Linux文件系统的核心操作指南
android·java·linux·开发语言·面试·系统架构
nightunderblackcat1 小时前
新手向:实现ATM模拟系统
java·开发语言·spring boot·spring cloud·tomcat·maven·intellij-idea
开开心心就好1 小时前
电脑息屏工具,一键黑屏超方便
开发语言·javascript·电脑·scala·erlang·perl
笑衬人心。1 小时前
Java 17 新特性笔记
java·开发语言·笔记
序属秋秋秋2 小时前
《C++初阶之内存管理》【内存分布 + operator new/delete + 定位new】
开发语言·c++·笔记·学习