mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-04-29 03:00:14 -04:00
32 lines
630 B
Python
32 lines
630 B
Python
import os
|
|
import time
|
|
import numpy as np
|
|
from efficientnet import EfficientNet
|
|
from tinygrad.tensor import Tensor
|
|
|
|
if __name__ == "__main__":
|
|
Tensor.default_gpu = os.getenv("GPU") is not None
|
|
model = EfficientNet()
|
|
|
|
BS = 4
|
|
|
|
img = np.zeros((BS,3,224,224), dtype=np.float32)
|
|
|
|
st = time.time()
|
|
out = model.forward(Tensor(img))
|
|
et = time.time()
|
|
print("forward %.2f s" % (et-st))
|
|
|
|
Y = [0]*BS
|
|
|
|
y = np.zeros((BS,1000), np.float32)
|
|
y[range(y.shape[0]),Y] = -1000.0
|
|
y = Tensor(y)
|
|
loss = out.logsoftmax().mul(y).mean()
|
|
|
|
st = time.time()
|
|
loss.backward()
|
|
et = time.time()
|
|
print("backward %.2f s" % (et-st))
|
|
|