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

柔光的暖阳◎ 2022-10-24 01:29 471阅读 0赞

# -*- 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();
    }