52 lines
1.7 KiB
Python
52 lines
1.7 KiB
Python
"""
|
|
NanoDetHeatV4 权重导出: PyTorch. pth → 二进制
|
|
"""
|
|
import sys, struct, os
|
|
sys.path.insert(0, os.path.dirname(__file__))
|
|
import torch
|
|
import numpy as np
|
|
from model_heat_v4 import NanoDetHeatV4
|
|
|
|
def export(pth_path, out_path):
|
|
model = NanoDetHeatV4(num_classes=3)
|
|
ckpt = torch.load(pth_path, map_location='cpu', weights_only=False)
|
|
if 'model' in ckpt:
|
|
ckpt = ckpt['model']
|
|
model.load_state_dict(ckpt, strict=False)
|
|
model.eval()
|
|
|
|
layers = {}
|
|
for name, param in model.named_parameters():
|
|
layers[name] = param.detach().cpu().numpy().astype(np.float32)
|
|
for name, buf in model.named_buffers():
|
|
layers[name] = buf.detach().cpu().numpy().astype(np.float32)
|
|
|
|
total = sum(v.size for v in layers.values())
|
|
print(f"层数: {len(layers)}, 总参数: {total:,} ≈ {total/1000:.1f}K")
|
|
|
|
for k, v in sorted(layers.items()):
|
|
print(f" {k:50s} {list(v.shape)} {v.size:,}")
|
|
|
|
with open(out_path, 'wb') as f:
|
|
f.write(struct.pack('I', len(layers)))
|
|
for name in sorted(layers.keys()):
|
|
data = layers[name]
|
|
nb = name.encode('utf-8')
|
|
f.write(struct.pack('I', len(nb)))
|
|
f.write(nb)
|
|
f.write(struct.pack('I', data.ndim))
|
|
f.write(struct.pack('i' * data.ndim, *data.shape))
|
|
f.write(data.tobytes())
|
|
|
|
print(f"\n导出: {out_path}")
|
|
|
|
# 测试推理
|
|
x = torch.randn(1, 3, 120, 160)
|
|
with torch.no_grad():
|
|
y = model(x)
|
|
print(f"输入: {x.shape} → 输出: {y.shape}")
|
|
|
|
if __name__ == '__main__':
|
|
base = os.path.dirname(os.path.abspath(__file__))
|
|
export(os.path.join(base, 'best.pth.1'), os.path.join(base, 'nanodet.bin'))
|