【pytorch基础】基于训练的pytorch模型转换为onnx模型并测试

前言

模型部署的过程中，不同的硬件可能支持不同的模型框架，本文介绍pytorch模型文件转换为onnx模型文件的实现过程，主要是基于Pytorch_Unet的实现过程，训练模型转换为onnx模型，并测试onnx的效果；

操作步骤

1. 基于训练完成的pth文件转换为onnx模型；

2. check和验证onnx模型；

3. 基于输入数据测试onnx模型；

实现过程

1. 基于训练完成的pth文件转换为onnx模型；

模型是基于Unet网络构建，基于Carvana数据集进行训练；

import io
import torch
import torch.onnx
from unet import UNet
import onnx
import onnxruntime
import numpy as np
from PIL import Image
import torchvision.transforms as transforms
from utils.dataset import BasicDataset

View Code