diff --git a/examples/hlb_cifar10.py b/examples/hlb_cifar10.py index f4802e2e36..692e7b0429 100644 --- a/examples/hlb_cifar10.py +++ b/examples/hlb_cifar10.py @@ -314,6 +314,9 @@ def train_cifar(): opt_bias = optim.SGD(params_bias, lr=0.01, momentum=hyp['opt']['momentum'], nesterov=True, weight_decay=hyp['opt']['bias_decay']) opt_non_bias = optim.SGD(params_non_bias, lr=0.01, momentum=hyp['opt']['momentum'], nesterov=True, weight_decay=hyp['opt']['non_bias_decay']) + # realize model params and optimizer state before JIT to avoid cache misses + Tensor.realize(*params_dict.values(), *opt_bias.b, *opt_non_bias.b) + # NOTE taken from the hlb_CIFAR repository, might need to be tuned initial_div_factor = hyp['opt']['initial_div_factor'] final_lr_ratio = hyp['opt']['final_lr_ratio'] diff --git a/tinygrad/nn/datasets.py b/tinygrad/nn/datasets.py index 30439fb49b..5c482208e2 100644 --- a/tinygrad/nn/datasets.py +++ b/tinygrad/nn/datasets.py @@ -8,7 +8,7 @@ def mnist(device=None, fashion=False): _mnist("t10k-images-idx3-ubyte.gz")[0x10:].reshape(-1,1,28,28).to(device), _mnist("t10k-labels-idx1-ubyte.gz")[8:].to(device) def cifar(device=None): - tt = tar_extract(Tensor.from_url('https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz', gunzip=True)) + tt = tar_extract(Tensor.from_url('https://data.brainchip.com/dataset-mirror/cifar10/cifar-10-binary.tar.gz', gunzip=True)) train = Tensor.cat(*[tt[f"cifar-10-batches-bin/data_batch_{i}.bin"].reshape(-1, 3073).to(device) for i in range(1,6)]) test = tt["cifar-10-batches-bin/test_batch.bin"].reshape(-1, 3073).to(device) return train[:, 1:].reshape(-1,3,32,32), train[:, 0], test[:, 1:].reshape(-1,3,32,32), test[:, 0]