#!/usr/bin/env python3
"""fp16 ONNX 第二法：权重转 half 后在 CUDA 上导出，包装层负责 float32 进出（onnxconverter_common 的 Cast 类型冲突绕开）。"""
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, 5, 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={k: {0: 'batch'} for k in ['views', 'scores', 'view', 'kind', 'zs', 'feat']}, 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')