github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-09 15:08:02 -05:00
Code Issues Packages Projects Releases Wiki Activity

cifar only shuffle data at epoch start (#3216)

Browse Source 
save 1ms CPU time per batch. also only shuffle training set
This commit is contained in:
chenyu
2024-01-23 14:41:22 -05:00
committed by GitHub
parent 4a07ea355d
commit 3c179cc27c
1 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
examples/hlb_cifar10.py
View File

@@ -185,8 +185,6 @@ def train_cifar():
while True:
st = time.monotonic()
X, Y = X_in, Y_in
order = list(range(0, X.shape[0]))
random.shuffle(order)
if is_train:
# TODO: these are not jitted
if getenv("RANDOM_CROP", 1):
@@ -196,14 +194,18 @@ def train_cifar():
if getenv("CUTMIX", 1):
if step >= hyp['net']['cutmix_steps']:
X, Y = cutmix(X, Y, mask_size=hyp['net']['cutmix_size'])
X, Y = X.numpy(), Y.numpy()
order = list(range(0, X.shape[0]))
random.shuffle(order)
X, Y = X.numpy()[order], Y.numpy()[order]
else:
X, Y = X.numpy(), Y.numpy()
et = time.monotonic()
print(f"shuffling {'training' if is_train else 'test'} dataset in {(et-st)*1e3:.2f} ms ({epoch=})")
for i in range(0, X.shape[0], BS):
# pad the last batch # TODO: not correct for test
batch_end = min(i+BS, Y.shape[0])
x = Tensor(X[order[batch_end-BS:batch_end],:])
y = Tensor(Y[order[batch_end-BS:batch_end]])
x = Tensor(X[batch_end-BS:batch_end], device=X_in.device)
y = Tensor(Y[batch_end-BS:batch_end], device=Y_in.device)
step += 1
yield x, y
epoch += 1