cifar flags for RANDOM_CROP, RANDOM_FLIP, and CUTMIX (#3204)

experimenting with different setups, also would like to jit the data augmentation next
This commit is contained in:
chenyu
2024-01-22 01:12:51 -05:00
committed by GitHub
parent 53afec2841
commit 99884f4c98

View File

@@ -205,10 +205,14 @@ def train_cifar():
order = list(range(0, X.shape[0]))
random.shuffle(order)
if is_train:
X = random_crop(X, crop_size=32)
X = Tensor.where(Tensor.rand(X.shape[0],1,1,1) < 0.5, X[..., ::-1], X) # flip LR
# NOTE: to bring cutmix back, make sure it's performing on mini-batch and not the whole set
# if step >= hyp['net']['cutmix_steps']: X, Y = cutmix(X, Y, mask_size=hyp['net']['cutmix_size'])
# TODO: these are not jitted
if getenv("RANDOM_CROP", 1):
X = random_crop(X, crop_size=32)
if getenv("RANDOM_FLIP", 1):
X = Tensor.where(Tensor.rand(X.shape[0],1,1,1) < 0.5, X[..., ::-1], X) # flip LR
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()
et = time.monotonic()
print(f"shuffling {'training' if is_train else 'test'} dataset in {(et-st)*1e3:.2f} ms ({cnt})")