mawei/tinygrad
1
0
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-06-13 00:15:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

train BS=16 for 32 steps

Browse Source
This commit is contained in:
George Hotz
2020-12-04 10:00:32 -08:00
parent ad1b225722
commit b8deb36e56
1 changed files with 17 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
examples/train_efficientnet.py
View File

@@ -8,24 +8,26 @@ if __name__ == "__main__":
Tensor.default_gpu = os.getenv("GPU") is not None
model = EfficientNet(int(os.getenv("NUM", "0")))
BS = 4
BS = 16
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))
for i in range(32):
print("running batch %d" % i)
Y = [0]*BS
st = time.time()
out = model.forward(Tensor(img))
et = time.time()
print("forward %.2f s" % (et-st))
y = np.zeros((BS,1000), np.float32)
y[range(y.shape[0]),Y] = -1000.0
y = Tensor(y)
loss = out.logsoftmax().mul(y).mean()
Y = [0]*BS
st = time.time()
loss.backward()
et = time.time()
print("backward %.2f s" % (et-st))
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))