improve benchmark_train_efficientnet
MODIFY
examples/benchmark_train_efficientnet.py
examples/benchmark_train_efficientnet.py
@@ -9,13 +9,21 @@ from tinygrad.tensor import Tensor
from test.test_gc import tensors_allocated
+try:
import pynvml
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+ def get_memory_used():
+ info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+ return info.used
+except Exception:
+ def get_memory_used():
+ return 0
NUM = int(os.getenv("NUM", 2))
BS = int(os.getenv("BS", 8))
CNT = int(os.getenv("CNT", 10))
+BACKWARD = int(os.getenv("BACKWARD", 0))
if __name__ == "__main__":
print(f"NUM:{NUM} BS:{BS} CNT:{CNT}")
@@ -31,16 +39,17 @@ if __name__ == "__main__":
st = time.monotonic()
out = model.forward(x_train)
loss = out.logsoftmax().mul(y_train).mean()
- optimizer.zero_grad()
+ if BACKWARD: optimizer.zero_grad()
loss.backward()
+ if BACKWARD: optimizer.step()
mt = time.monotonic()
- loss = loss.cpu().data[0]
+ loss.realize()
+ for p in parameters:
+ p.realize()
et = time.monotonic()
- info = pynvml.nvmlDeviceGetMemoryInfo(handle)
- print(f"{(et-st)*1000.0:7.2f} ms run, {(mt-st)*1000.0:7.2f} ms build, {loss:7.2f} loss, {tensors_allocated():4d} tensors, {info.used/1e9:.2f} GB used")
-
-
+ loss = loss.cpu().data[0]
+ print(f"{(et-st)*1000.0:7.2f} ms run, {(mt-st)*1000.0:7.2f} ms build, {loss:7.2f} loss, {tensors_allocated():4d} tensors, {get_memory_used()/1e9:.2f} GB used")

improve benchmark_train_efficientnet