用pytorch训练图像分类器模型导出ONNX测试_pytorch训练图像分类模型pth转onnx并测试_-牧野-的博客

onnx测试效果

网络很小，模型文件pth和cat_dog_classify.onnx大小只有63KB。通过OpenCV调用onnx，测试效果：

顺便贴一下py文件和cpp文件的代码（略长，文末有完整工程下载链接）

TrainTestConvertOnnx.py

# -*- coding: UTF-8 -*-
# Created by -牧野- CSDN https://blog.csdn.net/dcrmg/article/details/102807575
# 参考 https://blog.csdn.net/JR_Chan/article/details/95641758
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import time
import os
from math import ceil
import argparse
import copy
from PIL import Image
from torchvision import transforms, datasets
from torch.autograd import Variable
from tensorboardX import SummaryWriter
# 定义一个简单的二分类网络
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        # 三个卷积层用于提取特征
        # 1 input channel image 90x90, 8 output channel image 44x44
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=8, kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.MaxPool2d(2)
        # 8 input channel image 44x44, 16 output channel image 22x22
        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        # 16 input channel image 22x22, 32 output channel image 10x10
        self.conv3 = nn.Sequential(
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.MaxPool2d(2)
        self.classifier = nn.Sequential(
            nn.Linear(32 * 10 * 10, 3)
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = x.view(-1, 32 * 10 * 10)
        x = self.classifier(x)
        return x
# 训练模型入口
def train(args):
    # read data
    dataloders, dataset_sizes, class_names = ImageDataset(args)
    with open(args.class_file, 'w') as f:
        for name in class_names:
            f.writelines(name + '\n')
    # use gpu or not
    use_gpu = torch.cuda.is_available()
    print("use_gpu:{}".format(use_gpu))
    # get model
    model = SimpleNet()
    if args.resume:
        if os.path.isfile(args.resume):
            print(("=> loading checkpoint '{}'".format(args.resume)))
            model.load_state_dict(torch.load(args.resume))
        else:
            print(("=> no checkpoint found at '{}'".format(args.resume)))
    if use_gpu:
        model = torch.nn.DataParallel(model)
        model.to(torch.device('cuda'))
    else:
        model.to(torch.device('cpu'))
    # 用交叉熵损失函数(define loss function)
    criterion = nn.CrossEntropyLoss()
    # 梯度下降(Observe that all parameters are being optimized)
    optimizer_ft = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-4)
    # Decay LR by a factor of 0.98 every 1 epoch
    exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=1, gamma=0.98)
    model = train_model(args=args,
                        model=model,
                        criterion=criterion,
                        optimizer=optimizer_ft,
                        scheduler=exp_lr_scheduler,
                        num_epochs=args.num_epochs,
                        dataset_sizes=dataset_sizes,
                        use_gpu=use_gpu,
                        dataloders = dataloders)
    torch.save(model.state_dict(), os.path.join(args.save_path, 'best_model.pth'))
    writer.close()
# 测试单张图片（使用pth模型）入口
def test(test_model_path, test_img_path, class_file):
    best_model_path = test_model_path
    model = SimpleNet()
    model.load_state_dict(torch.load(best_model_path))
    model.eval()
    class_names = []
    with open(class_file, 'r') as f:
        lines = f.readlines()
        for line in lines:
            class_names.append(line)
    img_path = test_img_path
    predict_class = class_names[predict_image(model, img_path)]
    print(predict_class)
# 转换pytorch训练的pth模型到ONNX模型
def convert_model_to_ONNX(input_img_size, input_pth_model, output_ONNX):
    dummy_input = torch.randn(3, 1, input_img_size, input_img_size)
    model = SimpleNet()
    state_dict = torch.load(input_pth_model, map_location='cpu')
    model.load_state_dict(state_dict)
    model.eval()  # 设置模型为推理模式（重要）
    input_names = ["input_image"]
    output_names = ["output_classification"]
    torch.onnx.export(model, dummy_input, output_ONNX, verbose=True, input_names=input_names,
                      output_names=output_names)
# 训练模型主函数
def train_model(args, model, criterion, optimizer, scheduler, num_epochs, dataset_sizes, use_gpu, dataloders):
    begin = time.time()
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
    device = torch.device('cuda' if use_gpu else 'cpu')
    for epoch in range(args.start_epoch, num_epochs):
        # 每一个epoch中都有一个训练和一个验证过程(Each epoch has a training and validation phase)
        for phase in ['train', 'test']:
            if phase == 'train':
                scheduler.step(epoch)
                # 设置为训练模式(Set model to training mode)
                model.train()
            else:
                # 设置为验证模式(Set model to evaluate mode)
                model.eval()
            running_loss = 0.0
            running_corrects = 0
            tic_batch = time.time()
            # 在多个batch上依次处理数据(Iterate over data)
            for i, (inputs, labels) in enumerate(dataloders[phase]):
                inputs = inputs.to(device)
                labels = labels.to(device)
                # 梯度置零(zero the parameter gradients)
                optimizer.zero_grad()
                # 前向传播(forward)
                # 训练模式下才记录操作以进行反向传播(track history if only in train)
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)
                    # 训练模式下进行反向传播与梯度下降(backward + optimize only if in training phase)
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()
                # 统计损失和准确率(statistics)
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
                batch_loss = running_loss / (i * args.batch_size + inputs.size(0))
                batch_acc = running_corrects.double() / (i * args.batch_size + inputs.size(0))
                if phase == 'train' and (i + 1) % args.print_freq == 0:
                    print(
                        '[Epoch {}/{}]-[batch:{}/{}] lr:{:.6f} {} Loss: {:.6f}  Acc: {:.4f}  Time: {:.4f} sec/batch'.format(
                            epoch + 1, num_epochs, i + 1, ceil(dataset_sizes[phase] / args.batch_size),
                            scheduler.get_lr()[0], phase, batch_loss, batch_acc,
                            (time.time() - tic_batch) / args.print_freq))
                    tic_batch = time.time()
            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]
            if epoch == 0 and os.path.exists('result.txt'):
                os.remove('result.txt')
            with open('result.txt', 'a') as f:
                f.write('Epoch:{}/{} {} Loss: {:.4f} Acc: {:.4f} \n'.format(epoch + 1, num_epochs, phase, epoch_loss,
                                                                            epoch_acc))
            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
            writer.add_scalar(phase + '/Loss', epoch_loss, epoch)
            writer.add_scalar(phase + '/Acc', epoch_acc, epoch)
        if (epoch + 1) % args.save_epoch_freq == 0:
            if not os.path.exists(args.save_path):
                os.makedirs(args.save_path)
            torch.save(model.state_dict(), os.path.join(args.save_path, "epoch_" + str(epoch) + ".pth"))
        # 深拷贝模型(deep copy the model)
        if phase == 'test' and epoch_acc > best_acc:
            best_acc = epoch_acc
            best_model_wts = copy.deepcopy(model.state_dict())
    # 将model保存为graph
    writer.add_graph(model, (inputs,))
    time_elapsed = time.time() - begin
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
    print('Best val Accuracy: {:4f}'.format(best_acc))
    # 载入最佳模型参数(load best model weights)
    model.load_state_dict(best_model_wts)
    return model
# 测试单张图片主函数
def predict_image(model, image_path):
    image = Image.open(image_path).convert('L')
    # 测试时截取中间的90x90
    transformation1 = transforms.Compose([
        transforms.CenterCrop(90),
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5])
    # 预处理图像
    image_tensor = transformation1(image).float()
    # 额外添加一个批次维度，因为PyTorch将所有的图像当做批次
    image_tensor = image_tensor.unsqueeze_(0)
    if torch.cuda.is_available():
        image_tensor.cuda()
    # 将输入变为变量
    input = Variable(image_tensor)
    # 预测图像的类别
    output = model(input)
    index = output.data.numpy().argmax()
    return index
# 使用PIL读取图片并转换为灰度图
def readImg(path):
    im = Image.open(path)
    return im.convert("L")
# 读取训练和测试数据
def ImageDataset(args):
    # 数据增强及归一化
    # 图片都是120x120的，训练时随机裁取90x90的部分，测试时裁取中间的90x90
    data_transforms = {
        'train': transforms.Compose([
            transforms.RandomCrop(def_img_train_and_test_size),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5])
        'test': transforms.Compose([
            transforms.CenterCrop(def_img_train_and_test_size),
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5])
    data_dir = args.data_dir
    image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
                                            data_transforms[x], loader=readImg)
                    for x in ['train', 'test']}
    dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=args.batch_size,
                                                shuffle=(x == 'train'), num_workers=args.num_workers)
                for x in ['train', 'test']}
    dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'test']}
    class_names = image_datasets['train'].classes
    return dataloaders, dataset_sizes, class_names
# 设置参数
def set_parser():
    parser = argparse.ArgumentParser(description='classification')
    # 图片数据的根目录(Root catalog of images)
    parser.add_argument('--data-dir', type=str, default='images')
    parser.add_argument('--class-file', type=str, default='class_names.class')
    parser.add_argument('--batch-size', type=int, default=8)
    parser.add_argument('--num-epochs', type=int, default=100)
    parser.add_argument('--lr', type=float, default=0.002)  # those who set lr greater than 0.01 are hooligans!!
    parser.add_argument('--num-workers', type=int, default=8)
    parser.add_argument('--print-freq', type=int, default=100)
    parser.add_argument('--save-epoch-freq', type=int, default=1)
    parser.add_argument('--save-path', type=str, default='output')
    parser.add_argument('--resume', type=str, default='', help='For training from one checkpoint')
    parser.add_argument('--start-epoch', type=int, default=0, help='Corresponding to the epoch of resume')
    return parser.parse_args()
if __name__ == '__main__':
    writer = SummaryWriter(log_dir='log')
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    def_img_train_and_test_size = 90  # 训练尺寸
    args = set_parser()      # 设置参数
    train(args)      # 训练模型
    test('./output/best_model.pth', './images/test/cat/cat.0.jpg', args.class_file)     # 测试模型（单张图片）
    # 转换pytorch的pth模型到ONNX模型
    convert_model_to_ONNX(def_img_train_and_test_size, './output/epoch_99.pth', "./cat_dog_classify.onnx")

TestOnnx.cpp

// PthONNX.cpp : 基于OpenCV dnn、 onnx 的cat、dog二分类程序 // Created by -牧野- 2019年10月29日 https://blog.csdn.net/dcrmg/article/details/102807575 #include <iostream> #include <opencv2/highgui.hpp> #include <opencv2/imgproc.hpp> #include <opencv2/dnn.hpp> #include <fstream> //ONNX 执行推理类 class PthONNX { public: //@model_path ONNX模型路径 //@classes_file_path 分类信息文件 //@input_size 网络输入大小 PthONNX(const std::string &model_path, const std::string &classes_file_path, cv::Size input_size); //@input_image 输入图片，BGR格式 //@classification_output 网络输出的分类名称 0：cat 1：dog 1：None void Classify(const cv::Mat &input_image, std::string &classification_output); private: void ClassifyImplement(const cv::Mat &image, std::string &classification_output); private: cv::Size input_size_; cv::dnn::Net net_classify_; std::vector<std::string> classes_; // 构造函数 PthONNX::PthONNX(const std::string &model_path, const std::string &classes_file_path, cv::Size input_size) : input_size_(input_size) { std::ifstream ifs(classes_file_path.c_str()); assert(ifs.is_open()); std::string line; while (getline(ifs, line)) { line = line; classes_.push_back(line); net_classify_ = cv::dnn::readNetFromONNX(model_path); net_classify_.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV); net_classify_.setPreferableTarget(cv::dnn::DNN_TARGET_CPU); // ONNX推理入口函数 void PthONNX::Classify(const cv::Mat &input_image, std::string &classification_results) { assert(input_image.data); cv::Mat image = input_image.clone(); cv::resize(image, image, cv::Size(90, 90)); cv::cvtColor(image, image, cv::COLOR_BGR2GRAY); ClassifyImplement(image,classification_results); //ONNX推理主函数 void PthONNX::ClassifyImplement(const cv::Mat &image,std::string &classification_results) { classification_results.clear(); //***********前处理*********** cv::Scalar mean_value(0, 0, 0); cv::Mat input_blob = cv::dnn::blobFromImage(image, 1, input_size_, mean_value, false, false, CV_32F); //***********前处理*********** net_classify_.setInput(input_blob); const std::vector<cv::String> &out_names = net_classify_.getUnconnectedOutLayersNames(); cv::Mat out_tensor = net_classify_.forward(out_names[0]); //***********后处理*********** double minVal; double maxVal; cv::Point minIdx; cv::Point maxIdx; // minnimum Index, maximum Index cv::minMaxLoc(out_tensor, &minVal, &maxVal, &minIdx, &maxIdx); int index_class = maxIdx.x; classification_results = (index_class <= 1) ? classes_[index_class] : "None"; //***********后处理*********** int main() const std::string img_path = "D:/1/1/SimpleNet-master/images/train/cat/cat.4896.jpg"; const std::string onnx_model_path = "D:/1/1/pytorch-train-test-onnx/cat_dog_classify.onnx"; const std::string class_names_file_path = "D:/software/VS2019_Test/PthONNX/x64/class_names.class"; const cv::Size net_input_size(90, 90); cv::Mat img = cv::imread(img_path); std::string classify_output; // 分类结果 PthONNX classifier(onnx_model_path, class_names_file_path, net_input_size); classifier.Classify(img, classify_output); std::cout << "图片类别：" << classify_output << std::endl << std::endl; cv::putText(img, classify_output, cv::Point(20,20), 2, 1.2, cv::Scalar(0, 0, 255)); cv::imshow("classify", img); cv::waitKey();

完整工程（含数据集，pytorch训练和测试，pth模型转onnx，onnx文件加载和测试）下载链接： pytorch训练图像分类模型pth转ONNX并测试

很多时候有 pytorch 模型转 onnx 模型的必要，比如用tensorRT加速的时候。本文将介绍 pytorch 的 pth 模型如何转换成 onnx ，并且验证你转的模型对不对。先给官网链接：https:// pytorch .org/docs/stable/ onnx .html 咱们直接用一段代码来看：（本人亲自整理，有问题可留言交流~） import os.path as osp import numpy as np import onnx import onnx runtime as ort import t 要训练模型， main.py使用所需的模型架构和ImageNet数据集的路径运行main.py ： python main.py -a resnet18 [imagenet-folder with train and val folders] 默认学习率计划从0.1开始，每30个时代衰减10倍。这对于ResNet和具有批处理归一化的模型是合适的，但对于AlexNet和VGG来说太高了。使用0.01作为AlexNet或VGG的初始学习率： 1. 搭建自己的简单二分类网络，使用 pytorch 训练和测试； 2. 将 pytorch 训练的 pth 模型转换成 ONNX ，并编码测试； 3. 含训练和测试数据，含训练 ok的 pth 模型和 ONNX 模型，含完整python和C++实现； 4. 使用方法：首先运行“TrainTestConvert Onnx .py”执行“ 训练数据读入、模型训练、模型测试、导出 onnx ”，再运行“Test Onnx .cpp” 测试 onnx （需要配置 OpenCV ）；提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档 unet学习笔记（milesial/ Pytorch -UNet）原理讲解unet_parts.py内容predict.py内容所用到的代码 Pytorch -UNet/predict.py at master · milesial/ Pytorch -UNet ·GitHub 蓝/白色框表示 feature map；蓝色箭头表示 3x3 卷积，用于特征提取；灰色箭头表示skip-connection，用于特征融合；红色箭头. 输入图像的格式为[C, H, W]，即(channels, height, and width)，我们也需要提供一个batch size。batch size指一次处理多少张图像。所以输入图像格式为[N, C, H, W]。同时，图像的像素值要在0-1之间。 RetinaNet的输出格式它输出一个列表包括一个字典，其包含结果张量。格式为List[D... 文章目录1 准备工作1.1 efficientnet网络介绍1.2 efficientnet 训练分类数据集2 为何要转 3 安装相关依赖4 转换过程5 检验生成的 onnx 模型 5.1 onnx .checker检验5.2 np.testing.assert_allclose校验5.3 warning消除记录5.4 测试一张图片校验6 整合到一起的代码 1 准备工作 1.1 efficientnet网络介绍详见参考链接EfficientNet网络结构及代码详解。 1.2 efficientnet 训练分类数据集大部分的 pytorch 入门教程，都是使用torchvision里面的数据进行训练和测试。如果我们是自己的图片数据，又该怎么做呢？一、我的数据我在学习的时候，使用的是fashion-mnist。这个数据比较小，我的电脑没有GPU，还能吃得消。关于fashion-mnist数据，可以百度，也可以点此了解一下，数据就像这个样子：下载地址：https://github.com/zalandoresearch/fashion-mnist 但是下载下来是一种二进制文件，并不是图片，因此我先转换成了图片。我先解压gz文件到e:/fashion_mnist/文件夹然后运行代码： import from PIL import Image import torchvision.transforms as transforms import onnx runtime as rt import numpy as np ################################################################################################################. from onnx runtime.datasets import get_example import onnx runtime from onnx import shape_inference import onnx import os from models import * img_size = 416 cfg = 'cfg/yolov3.cfg'