github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-09 06:58:11 -05:00
Code Issues Packages Projects Releases Wiki Activity

yolov8 works up to the MaxPool

Browse Source
This commit is contained in:
George Hotz
2023-02-22 19:32:01 -08:00
parent 5190784cbb
commit 76b4d0577d
2 changed files with 21 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
examples/yolov8.py Normal file
View File

@@ -0,0 +1,17 @@
#!/usr/bin/env python3
import os
from ultralytics import YOLO
import onnx
from extra.onnx import get_run_onnx
from tinygrad.tensor import Tensor
os.chdir("/tmp")
if not os.path.isfile("yolov8n-seg.onnx"):
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 = get_run_onnx(onnx_model)
run_onnx({"images": Tensor.zeros(1,3,480,640)}, debug=True)

6
extra/onnx.py
View File

@@ -18,6 +18,7 @@ def get_run_onnx(onnx_model):
def attribute_parse(a):
if a.type == 7: return tuple([int(x) for x in a.ints])
elif a.type == 4: return buffer_parse(a.t) # TENSOR
elif a.type == 3: return str(a.s)
elif a.type == 2: return int(a.i)
elif a.type == 1: return float(a.f)
else: raise Exception(f"can't parse {a.type} {a}")
@@ -136,6 +137,7 @@ def get_run_onnx(onnx_model):
if n.op_type == "Sub": ret = inp[0] - inp[1]
if n.op_type == "Mul": ret = inp[0] * inp[1]
elif n.op_type == "Split":
if 'split' not in opt: opt['split'] = [int(x) for x in inp[1].numpy()] # split can be a tensor
i = 0
arg = [(0,x) for x in inp[0].shape]
for o,s in zip(n.output, opt['split']):
@@ -147,11 +149,11 @@ def get_run_onnx(onnx_model):
assert opt['kernel_shape'] == opt['strides'] or opt['strides'] == (1,1)
ret = inp[0].avg_pool2d(opt['kernel_shape'])
elif n.op_type == "MaxPool":
assert opt['kernel_shape'] == opt['strides']
assert opt['kernel_shape'] == opt['strides'], f"kernel_shape and stride mismatch {opt}"
#opt['kernel_shape'] = opt['strides']
# TODO: this is untested and probably wrong
ret = inp[0].pad2d(opt['pads'])
ret = ret.max_pool2d(opt['kernel_shape'])
ret = ret.max_pool2d(opt['kernel_shape'], opt['strides'])
# strides aren't supported in max_pool
#chan = ret.shape[1]
#w = Tensor.eye(chan).reshape((chan, chan, 1, 1))