NorgesGruppen-Data / scripts /export_fp16.py
Henrikalvados's picture
Upload folder using huggingface_hub
9f7ad84 verified
Raw
History Blame Contribute Delete
689 Bytes
import torch, onnx, os, sys
_orig = torch.load
def _safe(*a, **kw): kw["weights_only"] = False; return _orig(*a, **kw)
torch.load = _safe
from ultralytics import YOLO
from onnxconverter_common import float16
best = sys.argv[1] if len(sys.argv) > 1 else "runs/detect/train/weights/best.pt"
print(f"Exporting: {best}")
model = YOLO(best, task="detect")
model.export(format="onnx", opset=17, simplify=True, imgsz=1280)
onnx_path = best.replace(".pt", ".onnx")
m = onnx.load(onnx_path)
m16 = float16.convert_float_to_float16(m, keep_io_types=True)
onnx.save(m16, "/tmp/model_fp16.onnx")
size = os.path.getsize("/tmp/model_fp16.onnx") / 1024 / 1024
print(f"FP16: {size:.1f} MB")
print("DONE")