0041-pytorch入门-猫狗二分类-pth转onnx模型

# -*- 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")

onnx.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();
}