前言
模型部署的过程中,不同的硬件可能支持不同的模型框架,本文介绍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