export_tsstg_onnx.py

import torch
import numpy as np
from Actionsrecognition.Models import TwoStreamSpatialTemporalGraph
import os

def export_onnx(weight_file, onnx_file, device='cpu'):
    graph_args = {'strategy': 'spatial'}
    num_class = 7
    model = TwoStreamSpatialTemporalGraph(graph_args, num_class).to(device)
    model.load_state_dict(torch.load(weight_file, map_location=device))
    model.eval()

    # dummy input
    N, C1, T, V = 1, 3, 30, 14
    C2 = 2
    pts = torch.randn(N, C1, T, V, dtype=torch.float32)
    mot = torch.randn(N, C2, T-1, V, dtype=torch.float32)
    dummy_input = (pts, mot)

    # 注意这里用 (dummy_input,) 作为输入
    torch.onnx.export(
        model, (dummy_input,), onnx_file,
        input_names=['pts', 'mot'],
        output_names=['output'],
        dynamic_axes={'pts': {0: 'batch'}, 'mot': {0: 'batch'}, 'output': {0: 'batch'}},
        opset_version=12
    )
    print(f"Exported to {onnx_file}")

if __name__ == '__main__':
    weight_file = os.path.join(os.path.dirname(__file__), 'Models/TSSTG/tsstg-model.pth')
    onnx_file = os.path.join(os.path.dirname(__file__), 'Models/TSSTG/tsstg-model.onnx')
    export_onnx(weight_file, onnx_file)