前言
模型部署的过程中,不同的硬件可能支持不同的模型框架,本文介绍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
查看更多关于【pytorch教程】的文章