From b66361843ae714f5c57af96f0125f7802f824232 Mon Sep 17 00:00:00 2001 From: Alex Telon Date: Wed, 2 Aug 2023 02:16:10 +0200 Subject: [PATCH] Timing and Context can now be used as decorators (#1385) * Context and Timing can now be used as decorators * Using Timing decorator in quickstart.md The time formating is better and is a useful tool to learn. Old: Time: 3.5260659999912605 New: Time: 3526.14 ms * Updated env_vars documentation for Context * Added test for Context decorator * Put new import on same line as others --- docs/env_vars.md | 19 ++++++++++++-- docs/quickstart.md | 60 +++++++++++++++++++++----------------------- test/test_helpers.py | 9 +++++++ tinygrad/helpers.py | 6 ++--- 4 files changed, 58 insertions(+), 36 deletions(-) diff --git a/docs/env_vars.md b/docs/env_vars.md index 2dd04f5512..dfb56f1595 100644 --- a/docs/env_vars.md +++ b/docs/env_vars.md @@ -5,11 +5,26 @@ Most of these are self-explanatory, and are usually used to set an option at run Example: `GPU=1 DEBUG=4 python3 -m pytest` -The columns are: Variable, Possible Value(s) and Description. +However you can also decorate a function to set a value only inside that function. -- A `#` means that the variable can take any integer value. +```python +# in tensor.py (probably only useful if you are a tinygrad developer) +@Context(DEBUG=4) +def numpy(self) -> ... +``` + +Or use contextmanager to temporarily set a value inside some scope: + +```python +with Context(DEBUG=0): + a = Tensor.ones(10, 10) + a *= 2 +``` ## Global Variables +The columns of this list are are: Variable, Possible Value(s) and Description. + +- A `#` means that the variable can take any integer value. These control the behavior of core tinygrad even when used as a library. diff --git a/docs/quickstart.md b/docs/quickstart.md index ae06caef84..ac26c42cec 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -9,7 +9,7 @@ We need some imports to get started: ```python import numpy as np -import time +from tinygrad.helpers import Timing ``` ## Tensors @@ -221,23 +221,22 @@ We will be using the same batch size of 64 and will be evaluating for 1000 of th # set training flag to false Tensor.training = False -st = time.perf_counter() -avg_acc = 0 -for step in range(1000): - # random sample a batch - samp = np.random.randint(0, X_test.shape[0], size=(64)) - batch = Tensor(X_test[samp], requires_grad=False) - # get the corresponding labels - labels = Y_test[samp] +with Timing("Time: "): + avg_acc = 0 + for step in range(1000): + # random sample a batch + samp = np.random.randint(0, X_test.shape[0], size=(64)) + batch = Tensor(X_test[samp], requires_grad=False) + # get the corresponding labels + labels = Y_test[samp] - # forward pass - out = net(batch) + # forward pass + out = net(batch) - # calculate accuracy - pred = np.argmax(out.numpy(), axis=-1) - avg_acc += (pred == labels).mean() -print(f"Test Accuracy: {avg_acc / 1000}") -print(f"Time: {time.perf_counter() - st}") + # calculate accuracy + pred = np.argmax(out.numpy(), axis=-1) + avg_acc += (pred == labels).mean() + print(f"Test Accuracy: {avg_acc / 1000}") ``` ## And that's it @@ -266,23 +265,22 @@ from tinygrad.jit import TinyJit def jit(x): return net(x).realize() -st = time.perf_counter() -avg_acc = 0 -for step in range(1000): - # random sample a batch - samp = np.random.randint(0, X_test.shape[0], size=(64)) - batch = Tensor(X_test[samp], requires_grad=False) - # get the corresponding labels - labels = Y_test[samp] +with Timing("Time: "): + avg_acc = 0 + for step in range(1000): + # random sample a batch + samp = np.random.randint(0, X_test.shape[0], size=(64)) + batch = Tensor(X_test[samp], requires_grad=False) + # get the corresponding labels + labels = Y_test[samp] - # forward pass with jit - out = jit(batch) + # forward pass with jit + out = jit(batch) - # calculate accuracy - pred = np.argmax(out.numpy(), axis=-1) - avg_acc += (pred == labels).mean() -print(f"Test Accuracy: {avg_acc / 1000}") -print(f"Time: {time.perf_counter() - st}") + # calculate accuracy + pred = np.argmax(out.numpy(), axis=-1) + avg_acc += (pred == labels).mean() + print(f"Test Accuracy: {avg_acc / 1000}") ``` You will find that the evaluation time is much faster than before and that your accelerator utilization is much higher. diff --git a/test/test_helpers.py b/test/test_helpers.py index 7350eec1f7..62822f1555 100644 --- a/test/test_helpers.py +++ b/test/test_helpers.py @@ -90,5 +90,14 @@ with Context(VARIABLE=1): self.assertEqual(VARIABLE.value, 1) self.assertEqual(VARIABLE.value, 0) + def test_decorator(self): + @Context(VARIABLE=1, DEBUG=4) + def test(): + self.assertEqual(VARIABLE.value, 1) + + self.assertEqual(VARIABLE.value, 0) + test() + self.assertEqual(VARIABLE.value, 0) + if __name__ == '__main__': unittest.main() \ No newline at end of file diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 9d504688b8..ffca0eee16 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -1,5 +1,5 @@ from __future__ import annotations -import os, functools, platform, time, re +import os, functools, platform, time, re, contextlib from weakref import KeyedRef, ref from _weakref import _remove_dead_weakref # type: ignore import numpy as np @@ -26,7 +26,7 @@ def fromimport(mod, frm): return getattr(__import__(mod, fromlist=[frm]), frm) @functools.lru_cache(maxsize=None) def getenv(key, default=0): return type(default)(os.getenv(key, default)) -class Context: +class Context(contextlib.ContextDecorator): stack: ClassVar[List[dict[str, int]]] = [{}] def __init__(self, **kwargs): self.kwargs = kwargs def __enter__(self): @@ -52,7 +52,7 @@ class ContextVar: DEBUG, IMAGE = ContextVar("DEBUG", 0), ContextVar("IMAGE", 0) GRAPH, PRUNEGRAPH, GRAPHPATH = getenv("GRAPH", 0), getenv("PRUNEGRAPH", 0), getenv("GRAPHPATH", "/tmp/net") -class Timing(object): +class Timing(contextlib.ContextDecorator): def __init__(self, prefix="", on_exit=None, enabled=True): self.prefix, self.on_exit, self.enabled = prefix, on_exit, enabled def __enter__(self): self.st = time.perf_counter_ns() def __exit__(self, exc_type, exc_val, exc_tb):