mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-04-29 03:00:14 -04:00
dd82b4c913
* this * clean up * more clean ups and improve debug msg * more correct training toggler * remove manual training toggling * change some variable names * actually just add the training toggle for LIMIT envvar too * more refinement * __call__ and OnnxRunner * fix half pylint, other half is importing from onnx while this file is onnx.py, figure out later * ahhhh found another mistake * remove limit from __call__ --------- Co-authored-by: chenyu <chenyu@fastmail.com>
19 lines
637 B
Python
19 lines
637 B
Python
#!/usr/bin/env python3
|
|
import os
|
|
from ultralytics import YOLO
|
|
import onnx
|
|
from pathlib import Path
|
|
from extra.onnx import OnnxRunner
|
|
from tinygrad.tensor import Tensor
|
|
|
|
os.chdir("/tmp")
|
|
if not Path("yolov8n-seg.onnx").is_file():
|
|
model = YOLO("yolov8n-seg.pt")
|
|
model.export(format="onnx", imgsz=[480,640])
|
|
onnx_model = onnx.load(open("yolov8n-seg.onnx", "rb"))
|
|
# TODO: move get example inputs to onnx
|
|
input_shapes = {inp.name:tuple(x.dim_value for x in inp.type.tensor_type.shape.dim) for inp in onnx_model.graph.input}
|
|
print(input_shapes)
|
|
run_onnx = OnnxRunner(onnx_model)
|
|
run_onnx({"images": Tensor.zeros(1,3,480,640)}, debug=True)
|