onnx infer

[yxl502]

import torch import torchvision.models as models import sys import os current_dir = os.path.dirname(os.path.abspath(file))
parent_dir = os.path.dirname(current_dir) sys.path.insert(0, parent_dir) from tubevit.model import TubeViTLightningModule

加载一个预训练的PyTorch模型

model = TubeViTLightningModule( num_classes=3, video_shape=[3, 1, 448, 224], num_layers=12, num_heads=12, hidden_dim=768, mlp_dim=3072, weight_path="../weights/tubevit_vitbase_nc3_fpc3_448_224.pt", test_each_epoch = False ) model.eval()

定义模型的输入示例

dummy_input = torch.randn(1, 3, 1, 448, 224)

指定要保存的ONNX文件的路径

onnx_file_path = "./weights/test.onnx"

导出模型到ONNX格式

torch.onnx.export(model, dummy_input, onnx_file_path, verbose=True,opset_version=12) # 版本只有7-16，但是都不支持