PyTorch 训练图像分割网络(代码讲解)

整个工程文件已放到Github上
https://github.com/yaoyi30/PyTorch_Image_Segmentation

一、训练图像分类网络主要流程

1. 构建数据集
2. 数据预处理、包括数据增强和数据标准化和归一化
3. 构建网络模型
4. 设置学习率、优化器、损失函数等超参数
5. 训练和验证

二、各个流程简要说明

1. 构建数据集

本文使用supervisely 发布的人像分割数据集，百度网盘地址：

https://pan.baidu.com/s/1B8eBqg7XROHOsm5OLw-t9g 提取码: 52ss

在工程目录下，新建datasets文件夹，在文件夹内分别新建images和labels文件夹，用来放图片和对应的mask图片，之后在两个文件夹内新建train和val文件夹用来存放训练和验证数据，结构如下：

datasets/
  images/    # images
     train/
        img1.jpg
        img2.jpg
         .
         .
         .
     val/
        img1.jpg
        img2.jpg
         .
         .
         .
  labels/     # masks
     train/
        img1.png
        img2.png
         .
         .
         .
     val/
        img1.png
        img2.png
         .
         .
         .

2. 数据预处理

将图像resize到统一大小，之后转为tensor格式再进行标准化，预处理之后的图片可以正常输入网络，对于训练集可以采取一些数据增强手段来增强网络的泛化能力，验证集不做数据增强。

    #训练数据预处理、数据增强设置
    train_transform = Compose([
                                    Resize(args.input_size), #图像resize到统一大小
                                    RandomHorizontalFlip(0.5),  #数据增强，水平翻转
                                    ToTensor(), #转为tensor格式，值变为0-1之间
                                    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) #标准化
                             ])
    #验证数据预处理
    val_transform = Compose([
                                    Resize(args.input_size), #图像resize到统一大小
                                    ToTensor(), #转为tensor格式，值变为0-1之间
                                    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) #标准化
                            ])

3. 构建网络模型

本文搭建了一个简单的图像分割网路，命名为Simplify_Net。

    model = Simplify_Net(args.nb_classes)

4. 设置学习率、优化器、损失函数等超参数

    #定义损失函数，因为分割是像素级的分类，因此可以选用交叉熵损失函数
    loss_function = nn.CrossEntropyLoss()
    #定义优化器（初始学习率和权重衰减值）
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.init_lr, weight_decay=args.weight_decay)
    #定义学习率类型，此处选用余弦退火学习率，设置最大值
    scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, args.max_lr, total_steps=args.epochs, verbose=True)

5. 训练和验证

	#训练和验证模型，具体函写在了utils/engine.py文件中
	history = train_and_val(args.epochs, model, train_loader,val_loader,loss_function, optimizer,scheduler,args.output_dir,device,args.nb_classes) 

三、工程代码文件详细讲解

train.py

定义训练的入口函数，以及训练所需要的流程

1. 导入相应的库和文件

import os
import torch
import torch.nn as nn
from models.Simplify_Net import Simplify_Net
from utils.engine import train_and_val,plot_pix_acc,plot_miou,plot_loss,plot_lr
import argparse
import numpy as np
from utils.transform import Resize,Compose,ToTensor,Normalize,RandomHorizontalFlip
from utils.datasets import SegData

2. 训练参数设置

def get_args_parser():
    parser = argparse.ArgumentParser('Image Segmentation Train', add_help=False)
    parser.add_argument('--batch_size', default=32, type=int,help='Batch size for training')
    parser.add_argument('--epochs', default=100, type=int)
    parser.add_argument('--input_size', default=[224,224],nargs='+',type=int,help='images input size')
    parser.add_argument('--data_path', default='./datasets/', type=str,help='dataset path')

    parser.add_argument('--init_lr', default=1e-5, type=float,help='intial lr')
    parser.add_argument('--max_lr', default=1e-3, type=float,help='max lr')
    parser.add_argument('--weight_decay', default=1e-5, type=float,help='weight decay')

    parser.add_argument('--nb_classes', default=2, type=int,help='number of the classification types')
    parser.add_argument('--output_dir', default='./output_dir',help='path where to save, empty for no saving')
    parser.add_argument('--device', default='cuda',help='device to use for training / testing')
    parser.add_argument('--num_workers', default=4, type=int)

    return parser

3. 定义主函数

def main(args):

    device = torch.device(args.device)

    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    train_transform = Compose([
                                    Resize(args.input_size),
                                    RandomHorizontalFlip(0.5),
                                    ToTensor(),
                                    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
                             ])

    val_transform = Compose([
                                    Resize(args.input_size),
                                    ToTensor(),
                                    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
                            ])

    train_dataset = SegData(image_path=os.path.join(args.data_path, 'images/train'),
                            mask_path=os.path.join(args.data_path, 'labels/train'),
                            data_transforms=train_transform)
    train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=args.batch_size, shuffle=True,
                                               num_workers=args.num_workers)

    val_dataset = SegData(image_path=os.path.join(args.data_path, 'images/val'),
                            mask_path=os.path.join(args.data_path, 'labels/val'),
                            data_transforms=val_transform)
    val_loader = torch.utils.data.DataLoader(dataset=val_dataset, batch_size=args.batch_size, shuffle=False,
                                             num_workers=args.num_workers)

    model = Simplify_Net(args.nb_classes)
    loss_function = nn.CrossEntropyLoss()

    optimizer = torch.optim.AdamW(model.parameters(), lr=args.init_lr, weight_decay=args.weight_decay)
    scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, args.max_lr, total_steps=args.epochs, verbose=True)

    history = train_and_val(args.epochs, model, train_loader,val_loader,loss_function, optimizer,scheduler,args.output_dir,device,args.nb_classes)

    plot_loss(np.arange(0,args.epochs),args.output_dir, history)
    plot_pix_acc(np.arange(0,args.epochs),args.output_dir, history)
    plot_miou(np.arange(0,args.epochs),args.output_dir, history)
    plot_lr(np.arange(0,args.epochs),args.output_dir, history)

4. 开始执行

if __name__ == '__main__':
    #获取训练参数
    args = get_args_parser()
    #解析训练参数
    args = args.parse_args()
    #训练参数传入主函数
    main(args)

运行train.py，训练时打印的信息，包括每一轮的学习率，训练集和验证集指标，运行时间等

Simplify_Net.py

定义网络结构，本文定义一个简单的Encoder-Decoder结构的卷积神经网络

import torch
import torch.nn as nn
from torch.nn.functional import interpolate

class Simplify_Net(nn.Module):
    def __init__(self, num_classes=2):
        super(Simplify_Net, self).__init__()

        self.conv1 = nn.Conv2d(in_channels=3,out_channels=16,kernel_size=3,padding=1,stride=2)
        self.bn1 = nn.BatchNorm2d(16)
        self.relu1 = nn.ReLU(inplace=True)

        self.conv2 = nn.Conv2d(in_channels=16,out_channels=16,kernel_size=3,padding=1,stride=2)
        self.bn2 = nn.BatchNorm2d(16)
        self.relu2 = nn.ReLU(inplace=True)

        self.conv3 = nn.Conv2d(in_channels=16,out_channels=16,kernel_size=3,padding=1,stride=2)
        self.bn3 = nn.BatchNorm2d(16)
        self.relu3 = nn.ReLU(inplace=True)

        self.upconv1 = nn.ConvTranspose2d(in_channels=16,out_channels=16,kernel_size=4,padding=1,stride=2)
        self.bn4 = nn.BatchNorm2d(16)
        self.relu4 = nn.ReLU(inplace=True)

        self.upconv2 = nn.ConvTranspose2d(in_channels=32,out_channels=16,kernel_size=4,padding=1,stride=2)
        self.bn5 = nn.BatchNorm2d(16)
        self.relu5 = nn.ReLU(inplace=True)

        self.conv_last = nn.Conv2d(in_channels=32,out_channels=num_classes,kernel_size=1,stride=1)


    def forward(self, x):

        x1 = self.relu1(self.bn1(self.conv1(x)))
        x2 = self.relu2(self.bn2(self.conv2(x1)))
        x3 = self.relu3(self.bn3(self.conv3(x2)))

        up1 = torch.cat([x2,self.relu4(self.bn4(self.upconv1(x3)))],dim=1)
        up2 = torch.cat([x1,self.relu5(self.bn5(self.upconv2(up1)))],dim=1)

        up3 = self.conv_last(up2)

        out = interpolate(up3, scale_factor=2, mode='bilinear', align_corners=False)

        return out

utils/datasets.py

定义数据读取的类

import os
from torch.utils.data import Dataset
from PIL import Image


class SegData(Dataset):
    def __init__(self, image_path, mask_path, data_transforms=None):
        self.image_path = image_path
        self.mask_path = mask_path

        self.images = os.listdir(self.image_path)
        self.masks = os.listdir(self.mask_path)
        self.transform = data_transforms

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        image_filename = self.images[idx]
        mask_filename = image_filename.replace('jpeg','png')

        image = Image.open(os.path.join(self.image_path, image_filename)).convert('RGB')
        mask = Image.open(os.path.join(self.mask_path, mask_filename)).convert('L')

        if self.transform is not None:
            image, mask = self.transform(image ,mask)

        return image, mask

utils/transform.py

定义数据预处理的类

import numpy as np
import random
import torch
from torchvision.transforms import functional as F

# 将img和mask resize到统一大小
class Resize(object):
    def __init__(self, size):
        self.size = size

    def __call__(self, image, target=None):
        image = F.resize(image, self.size)
        if target is not None:
            target = F.resize(target, self.size, interpolation=F.InterpolationMode.NEAREST)
        return image, target

#随机左右翻转
class RandomHorizontalFlip(object):
    def __init__(self, flip_prob):
        self.flip_prob = flip_prob

    def __call__(self, image, target=None):
        if random.random() < self.flip_prob:
            image = F.hflip(image)
            if target is not None:
                target = F.hflip(target)
        return image, target

#标准化
class Normalize(object):
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, image, target):
        image = F.normalize(image, mean=self.mean, std=self.std)
        return image, target

#img转tensor,值变为0-1之间,label直接转为tensor
class ToTensor(object):
    def __call__(self, image, target):
        image = F.to_tensor(image)
        if target is not None:
            target = torch.as_tensor(np.array(target), dtype=torch.int64)
        return image, target

#不同数据预处理类组合起来
class Compose(object):
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, image, mask=None):
        for t in self.transforms:
            image, mask = t(image, mask)
        return image, mask

utils/metrics.py

定义计算像素准确率、MIoU等指标的类

import numpy as np

class Evaluator(object):
    def __init__(self, num_class):
        self.num_class = num_class
        self.confusion_matrix = np.zeros((self.num_class,)*2)
    #计算像素准确率
    def Pixel_Accuracy(self):
        Acc = np.diag(self.confusion_matrix).sum() / self.confusion_matrix.sum()
        return Acc

    def Pixel_Accuracy_Class(self):
        Acc = np.diag(self.confusion_matrix) / self.confusion_matrix.sum(axis=1)
        Acc = np.nanmean(Acc)
        return Acc
    #计算每一类IoU和MIoU
    def Mean_Intersection_over_Union(self):
        IoU = np.diag(self.confusion_matrix) / (
                    np.sum(self.confusion_matrix, axis=1) + np.sum(self.confusion_matrix, axis=0) -
                    np.diag(self.confusion_matrix))
        MIoU = np.nanmean(IoU)
        return IoU,MIoU

    def Frequency_Weighted_Intersection_over_Union(self):
        freq = np.sum(self.confusion_matrix, axis=1) / np.sum(self.confusion_matrix)
        iu = np.diag(self.confusion_matrix) / (
                    np.sum(self.confusion_matrix, axis=1) + np.sum(self.confusion_matrix, axis=0) -
                    np.diag(self.confusion_matrix))

        FWIoU = (freq[freq > 0] * iu[freq > 0]).sum()
        return FWIoU

    def _generate_matrix(self, gt_image, pre_image):
        mask = (gt_image >= 0) & (gt_image < self.num_class)
        label = self.num_class * gt_image[mask].astype('int') + pre_image[mask]
        count = np.bincount(label, minlength=self.num_class**2)
        confusion_matrix = count.reshape(self.num_class, self.num_class)
        return confusion_matrix
    #加入数据
    def add_batch(self, gt_image, pre_image):
        assert gt_image.shape == pre_image.shape
        self.confusion_matrix += self._generate_matrix(gt_image, pre_image)
    #重置
    def reset(self):
        self.confusion_matrix = np.zeros((self.num_class,) * 2)

utils/engine.py

定义具体的训练、验证以及绘制指标曲线的函数

1. 导入相应的库和文件

import os
import torch
import time
from tqdm import tqdm
import matplotlib.pyplot as plt
from utils.metrics import Evaluator
import numpy as np

2. 训练验证函数

def train_and_val(epochs, model, train_loader, val_loader,criterion, optimizer,scheduler,output_dir,device,nb_classes):

    train_loss = []
    val_loss = []
    train_pix_acc = []
    val_pix_acc = []
    train_miou = []
    val_miou = []
    learning_rate = []
    best_miou = 0

    segmetric_train = Evaluator(nb_classes)
    segmetric_val = Evaluator(nb_classes)

    model.to(device)

    fit_time = time.time()
    for e in range(epochs):

        torch.cuda.empty_cache()
        segmetric_train.reset()
        segmetric_val.reset()

        since = time.time()
        training_loss = 0

        model.train()
        with tqdm(total=len(train_loader)) as pbar:
            for image, label in train_loader:

                image = image.to(device)
                label = label.to(device)

                output = model(image)
                loss = criterion(output, label)

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                pred = output.data.cpu().numpy()
                label = label.cpu().numpy()
                pred = np.argmax(pred, axis=1)

                training_loss += loss.item()
                segmetric_train.add_batch(label, pred)
                pbar.update(1)

        model.eval()
        validation_loss = 0

        with torch.no_grad():
            with tqdm(total=len(val_loader)) as pb:
                for image, label in val_loader:

                    image = image.to(device)
                    label = label.to(device)

                    output = model(image)
                    loss = criterion(output, label)

                    pred = output.data.cpu().numpy()
                    label = label.cpu().numpy()
                    pred = np.argmax(pred, axis=1)

                    validation_loss += loss.item()
                    segmetric_val.add_batch(label, pred)
                    pb.update(1)

        train_loss.append(training_loss / len(train_loader))
        val_loss.append(validation_loss / len(val_loader))

        train_pix_acc.append(segmetric_train.Pixel_Accuracy())
        val_pix_acc.append(segmetric_val.Pixel_Accuracy())

        train_miou.append(segmetric_train.Mean_Intersection_over_Union()[1])
        val_miou.append(segmetric_val.Mean_Intersection_over_Union()[1])

        learning_rate.append(scheduler.get_last_lr())

        torch.save(model.state_dict(), os.path.join(output_dir,'last.pth'))
        if best_miou < segmetric_val.Mean_Intersection_over_Union()[1]:
            torch.save(model.state_dict(), os.path.join(output_dir,'best.pth'))


        print("Epoch:{}/{}..".format(e + 1, epochs),
              "Train Pix Acc: {:.3f}".format(segmetric_train.Pixel_Accuracy()),
              "Val Pix Acc: {:.3f}".format(segmetric_val.Pixel_Accuracy()),
              "Train MIoU: {:.3f}".format(segmetric_train.Mean_Intersection_over_Union()[1]),
              "Val MIoU: {:.3f}".format(segmetric_val.Mean_Intersection_over_Union()[1]),
              "Train Loss: {:.3f}".format(training_loss / len(train_loader)),
              "Val Loss: {:.3f}".format(validation_loss / len(val_loader)),
              "Time: {:.2f}s".format((time.time() - since)))

        scheduler.step()

    history = {'train_loss': train_loss, 'val_loss': val_loss ,'train_pix_acc': train_pix_acc, 'val_pix_acc': val_pix_acc,'train_miou': train_miou, 'val_miou': val_miou,'lr':learning_rate}
    print('Total time: {:.2f} m'.format((time.time() - fit_time) / 60))

    return history

3. 打印损失值曲线

def plot_loss(x,output_dir, history):
    plt.plot(x, history['val_loss'], label='val', marker='o')
    plt.plot(x, history['train_loss'], label='train', marker='o')
    plt.title('Loss per epoch')
    plt.ylabel('loss')
    plt.xlabel('epoch')
    plt.legend(), plt.grid()
    plt.savefig(os.path.join(output_dir,'loss.png'))
    plt.clf()

4. 打印像素准确率曲线

def plot_pix_acc(x,output_dir, history):
    plt.plot(x, history['train_pix_acc'], label='train_pix_acc', marker='x')
    plt.plot(x, history['val_pix_acc'], label='val_pix_acc', marker='x')
    plt.title('Pix Acc per epoch')
    plt.ylabel('pixal accuracy')
    plt.xlabel('epoch')
    plt.legend(), plt.grid()
    plt.savefig(os.path.join(output_dir,'pix_acc.png'))
    plt.clf()

网络结构较为简单，因此像素准确率不是特别的高

5. 打印MIoU曲线

def plot_miou(x,output_dir, history):
    plt.plot(x, history['train_miou'], label='train_miou', marker='x')
    plt.plot(x, history['val_miou'], label='val_miou', marker='x')
    plt.title('MIoU per epoch')
    plt.ylabel('miou')
    plt.xlabel('epoch')
    plt.legend(), plt.grid()
    plt.savefig(os.path.join(output_dir,'miou.png'))
    plt.clf()

网络结构较为简单，因此像MIoU不是特别的高

6. 打印学习率曲线

def plot_lr(x,output_dir,  history):
    plt.plot(x, history['lr'], label='learning_rate', marker='x')
    plt.title('learning rate per epoch')
    plt.ylabel('Learning_rate')
    plt.xlabel('epoch')
    plt.legend(), plt.grid()
    plt.savefig(os.path.join(output_dir,'learning_rate.png'))
    plt.clf()

从学习率曲线可以看出，约前30轮为warmup阶段，最大学习率为0.001

predict.py

进行单张图片预测

1. 导入相应的库和文件

import argparse
import torch
import torch.nn as nn
import torchvision.transforms as T
from models.Simplify_Net import Simplify_Net
from PIL import Image

2. 单张预测参数设置

def get_args_parser():
    parser = argparse.ArgumentParser('Predict Image', add_help=False)
    parser.add_argument('--image_path', default='./people-man-model-glasses-46219.jpeg', type=str, metavar='MODEL',help='Name of model to train')
    parser.add_argument('--input_size', default=[224,224],nargs='+',type=int,help='images input size')
    parser.add_argument('--weights', default='./output_dir/last.pth', type=str,help='dataset path')
    parser.add_argument('--nb_classes', default=2, type=int,help='number of the classification types')
    parser.add_argument('--device', default='cuda',help='device to use for training / testing')

    return parser

3. 定义主函数

def main(args):
    device = torch.device(args.device)

    image = Image.open(args.image_path).convert('RGB')
    img_size = image.size

    transforms = T.Compose([
        T.Resize(args.input_size),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225]),
    ])

    model = Simplify_Net(args.nb_classes)

    checkpoint = torch.load(args.weights, map_location='cpu')
    msg = model.load_state_dict(checkpoint, strict=True)
    print(msg)

    model.to(device)
    model.eval()

    input_tensor = transforms(image).unsqueeze(0).to(device)
    with torch.no_grad():
        output = model(input_tensor)
        pred = output.argmax(1).squeeze(0).cpu().numpy().astype(np.uint8)

    mask = Image.fromarray(pred)
    out = mask.resize(img_size)
    out.save("result.png")

4. 开始执行

if __name__ == '__main__':
    #获取训练参数
    args = get_args_parser()
    #解析训练参数
    args = args.parse_args()
    #训练参数传入主函数
    main(args)

运行predict.py，保存模型预测的结果

eval.py

进行模型整体指标评价

1. 导入相应的库和文件

import argparse
from utils.transform import Resize,Compose,ToTensor,Normalize,RandomHorizontalFlip
from utils.datasets import SegData
import torch
import os
import numpy as np
from tqdm import tqdm
from models.Simplify_Net import Simplify_Net
from utils.metrics import Evaluator

2. 模型评价参数设置

def get_args_parser():
    parser = argparse.ArgumentParser('Eval Model', add_help=False)
    parser.add_argument('--batch_size', default=1, type=int,help='Batch size for training')
    parser.add_argument('--input_size', default=[224,224],nargs='+',type=int,help='images input size')
    parser.add_argument('--data_path', default='./datasets/', type=str,help='dataset path')
    parser.add_argument('--weights', default='./output_dir/best.pth', type=str,help='dataset path')
    parser.add_argument('--nb_classes', default=2, type=int,help='number of the classification types')
    parser.add_argument('--device', default='cuda',help='device to use for training / testing')
    parser.add_argument('--num_workers', default=4, type=int)

    return parser

3. 定义主函数

def main(args):

    device = torch.device(args.device)

    segmetric = Evaluator(args.nb_classes)
    segmetric.reset()

    val_transform = Compose([
                                    Resize(args.input_size),
                                    ToTensor(),
                                    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
                            ])

    val_dataset = SegData(image_path=os.path.join(args.data_path, 'images/val'),
                            mask_path=os.path.join(args.data_path, 'labels/val'),
                            data_transforms=val_transform)
    val_loader = torch.utils.data.DataLoader(dataset=val_dataset, batch_size=args.batch_size, shuffle=False,
                                             num_workers=args.num_workers)

    model = Simplify_Net(args.nb_classes)

    checkpoint = torch.load(args.weights, map_location='cpu')
    msg = model.load_state_dict(checkpoint, strict=True)
    print(msg)

    model.to(device)
    model.eval()

    classes = ["background","human"]

    with torch.no_grad():
        with tqdm(total=len(val_loader)) as pbar:
            for image, label in val_loader:
                output = model(image.to(device))
                pred = output.data.cpu().numpy()
                label = label.cpu().numpy()
                pred = np.argmax(pred, axis=1)
                segmetric.add_batch(label, pred)
                pbar.update(1)

    pix_acc = segmetric.Pixel_Accuracy()
    every_iou,miou = segmetric.Mean_Intersection_over_Union()

    print("Pixel Accuracy is :", pix_acc)
    print("==========Every IOU==========")
    for name,prob in zip(classes,every_iou):
        print(name+" : "+str(prob))
    print("=============================")
    print("MiOU is :", miou)

4. 开始执行

if __name__ == '__main__':
    #获取训练参数
    args = get_args_parser()
    #解析训练参数
    args = args.parse_args()
    #训练参数传入主函数
    main(args)

运行eval.py，打印模型在验证集上的像素准确率，MIoU值和每一类的IoU值

export_onnx.py

将训练好的模型转onnx格式，以进行后续应用

1. 导入相应的库和文件

import torch
from models.Simplify_Net import Simplify_Net
import argparse

2. 转onnx模型参数设置

def get_args_parser():
    parser = argparse.ArgumentParser('Export Onnx', add_help=False)
    parser.add_argument('--input_size', default=[224,224],nargs='+',type=int,help='images input size')
    parser.add_argument('--weights', default='./output_dir/best.pth', type=str,help='dataset path')
    parser.add_argument('--nb_classes', default=2, type=int,help='number of the classification types')

    return parser

3. 定义主函数

def main(args):

    x = torch.randn(1, 3, args.input_size[0],args.input_size[1])
    input_names = ["input"]
    out_names = ["output"]

    model = Simplify_Net(args.nb_classes)

    checkpoint = torch.load(args.weights, map_location='cpu')
    msg = model.load_state_dict(checkpoint, strict=True)
    print(msg)

    model.eval()

    torch.onnx.export(model, x, args.weights.replace('pth','onnx'), export_params=True, training=False, input_names=input_names, output_names=out_names)
    print('please run: python -m onnxsim test.onnx test_sim.onnx\n')

4. 开始执行

if __name__ == '__main__':
    #获取训练参数
    args = get_args_parser()
    #解析训练参数
    args = args.parse_args()
    #训练参数传入主函数
    main(args)

运行export_onnx.py，之后进行模型的简化

简化之前（左）和之后（右）的onnx模型结构对比