#!/usr/bin/env python3
"""v3 fp16 ONNX：权重转 half 后在 CUDA 上导出，包装层负责 float32 进出。输入 views [B,V,3,224,224]，V 可为 5（整图＋四格）或 10（加滑窗），视角 0 须为整图。"""
import sys, os, torch, torch.nn as nn, onnx
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))); from shixiang import Shixiang
m = Shixiang('.', 'cpu')
class Wrap(nn.Module):
    def __init__(s, net): super().__init__(); s.net = net
    def forward(s, x):
        f, lo, kd, zs = s.net(x.half()); P = torch.sigmoid(lo.float()); sc, vw = P.max(1)
        return sc, vw.to(torch.int32), torch.softmax(kd.float(), -1), zs.float(), f[:, 0].float()
w = Wrap(m.net).half().cuda().eval(); x = torch.randn(1, 10, 3, 224, 224, device='cuda')
torch.onnx.export(w, x, 'shixiang_fp16.onnx', input_names=['views'], output_names=['scores', 'view', 'kind', 'zs', 'feat'],
                  dynamic_axes=dict({k: {0: 'batch'} for k in ['scores', 'view', 'kind', 'zs', 'feat']}, views={0: 'batch', 1: 'nviews'}), opset_version=17, dynamo=False, do_constant_folding=True)
mo = onnx.load('shixiang_fp16.onnx'); onnx.checker.check_model(mo); print('fp16 ok', os.path.getsize('shixiang_fp16.onnx') / 1e6, 'MB')