mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-01-24 06:18:01 -05:00
59 lines
1.5 KiB
Python
59 lines
1.5 KiB
Python
#!/usr/bin/env python3
|
|
import os
|
|
import time
|
|
from tqdm import trange
|
|
from extra.utils import get_parameters
|
|
from models.efficientnet import EfficientNet
|
|
import tinygrad.optim as optim
|
|
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}")
|
|
model = EfficientNet(NUM, classes=1000, has_se=False, track_running_stats=False)
|
|
parameters = get_parameters(model)
|
|
optimizer = optim.SGD(parameters, lr=0.001)
|
|
|
|
Tensor.training = True
|
|
for i in trange(CNT):
|
|
x_train = Tensor.randn(BS, 3, 224, 224, requires_grad=False)
|
|
y_train = Tensor.randn(BS, 1000, requires_grad=False)
|
|
|
|
st = time.monotonic()
|
|
out = model.forward(x_train)
|
|
loss = out.logsoftmax().mul(y_train).mean()
|
|
if BACKWARD:
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
mt = time.monotonic()
|
|
loss.realize()
|
|
for p in parameters:
|
|
p.realize()
|
|
et = time.monotonic()
|
|
|
|
loss = loss.detach().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")
|
|
|
|
|
|
|
|
|
|
|