github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-07 03:00:26 -04:00
Code Issues Packages Projects Releases Wiki Activity

run on update_benchmark too (#1460)

Browse Source 
* run on update_benchmark too

* amd inference test

* name it better

* add 10 CIFAR training steps
This commit is contained in:
George Hotz
2023-08-06 08:58:37 -07:00
committed by GitHub
parent 3d41674b42
commit 2ab282bfec
2 changed files with 28 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
test/external/external_model_benchmark.py vendored
View File

@@ -3,9 +3,11 @@ import pathlib
import time
import onnx
import torch
torch.set_num_threads(1)
from onnx2torch import convert
from extra.utils import download_file
from extra.onnx import get_run_onnx
from tinygrad.helpers import OSX
from tinygrad.tensor import Tensor
from tinygrad.lazy import Device
@@ -53,7 +55,7 @@ def benchmark_model(m):
np_inputs = {k:torch.randn(shp).numpy() for k,shp in input_shapes.items()}
assert len(input_shapes) < 20
for device in ["METAL", "CLANG"]:
for device in ["METAL" if OSX else "GPU", "CLANG"]:
Device.DEFAULT = device
inputs = {k:Tensor(inp) for k,inp in np_inputs.items()}
tinygrad_model = get_run_onnx(onnx_model)
@@ -69,9 +71,10 @@ def benchmark_model(m):
torch_inputs = [torch.tensor(x) for x in np_inputs.values()]
benchmark(m, "torch_cpu", lambda: torch_model(*torch_inputs))
torch_mps_model = torch_model.to('mps')
torch_mps_inputs = [x.to('mps') for x in torch_inputs]
benchmark(m, "torch_mps", lambda: torch_mps_model(*torch_mps_inputs))
torch_device = "mps" if OSX else "cuda"
torch_mps_model = torch_model.to(torch_device)
torch_mps_inputs = [x.to(torch_device) for x in torch_inputs]
benchmark(m, f"torch_{torch_device}", lambda: torch_mps_model(*torch_mps_inputs))
if open_csv is None:
open_csv = csv.DictWriter(open('onnx_inference_speed.csv', 'w', newline=''), fieldnames=list(CSV.keys()))