pth模型转onnx模型记录

原创

已于 2022-09-05 23:09:19 修改 · 1.6k 阅读

13 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #python #人工智能

于 2022-09-05 19:20:13 首次发布

本文详细记录了如何将使用nn.Sequential的PyTorch模型拆分为特征提取网络、图像特征网络和head部分，并分别导出为ONNX格式，以便后续进行RKNN转换，同时提供了关键代码和步骤。

记录一次pth转onnx

pth转onnx：

nanotrack中pth转onnx（原项目）在这里插入图片描述
因为使用nn.Sequential()定义网络，只接受单输入单输出，所以要将原模型进行解构：模板T的特征提取网络、图像X的特征提取网络、模型head部分。这三部分均要进行保存、转换过程。（转换过程中设置好输入参数，方便进行下一步rknn的转换）

pth转onnx：

import argparse
import os
import torch
import sys 

sys.path.append(os.getcwd()) 

from nanotrack.core.config import cfg

from nanotrack.utils.model_load import load_pretrain
from nanotrack.models.model_builder import ModelBuilder

os.environ["CUDA_VISIBLE_DEVICES"] = "0"  

parser = argparse.ArgumentParser(description='lighttrack')

parser.add_argument('--config', type=str, default='./m