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

touchup resnet_layer_bench (#4191)

Browse Source
This commit is contained in:
David Hou
2024-04-16 11:43:00 -07:00
committed by GitHub
parent 1dbf3b2b19
commit 7fb220a567
1 changed files with 3 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
test/external/external_benchmark_resnet.py vendored
View File

@@ -2,8 +2,8 @@ import functools
import time
import unittest
from tinygrad import Tensor, TinyJit
from tinygrad.helpers import getenv, GlobalCounters
from tinygrad import Tensor, TinyJit, GlobalCounters
from tinygrad.helpers import getenv
from tinygrad.nn.optim import SGD
from tinygrad.nn.state import get_parameters
@@ -82,22 +82,16 @@ class BenchmarkResnetTrain(unittest.TestCase):
mem_used = GlobalCounters.mem_used # a little high with JITCNT > 1 fsr
kernels = GlobalCounters.kernel_count // JITCNT
tm = (et-st) / JITCNT
best_tm = tm if best_tm is None or tm < best_tm else best_tm
if best_tm is None or tm < best_tm: best_tm = tm
print(f"\r{name:42s}: {best_tm * 1000:>9.2f} ms, {flops / 10**12 / tm:>7.2f} tflops, {mem_used / 10**9: 7.2f} GB used, {kernels:>6d} kernels")
def test_layer1_1(self): self._test_layer(*self._get_layer(0, 0))
def test_layer1_2(self): self._test_layer(*self._get_layer(0, 1))
def test_layer2_1(self): self._test_layer(*self._get_layer(1, 0))
def test_layer2_2(self): self._test_layer(*self._get_layer(1, 1))
def test_layer3_1(self): self._test_layer(*self._get_layer(2, 0))
def test_layer3_2(self): self._test_layer(*self._get_layer(2, 1))
def test_layer4_1(self): self._test_layer(*self._get_layer(3, 0))
def test_layer4_2(self): self._test_layer(*self._get_layer(3, 1))
if __name__ == '__main__':