Files
MoreTore/tinygrad_repo/extra/datasets/preprocess_imagenet.py
T
FrogAi 81da8ac9f4 openpilot v0.9.7 release
date: 2024-06-11T01:36:39
master commit: f8cb04e4a8b032b72a909f68b808a50936184bee
2024-07-24 10:32:22 -07:00

23 lines
696 B
Python

from tinygrad.helpers import dtypes
from tinygrad.tensor import Tensor
from extra.datasets.imagenet import iterate, get_val_files
if __name__ == "__main__":
#sz = len(get_val_files())
sz = 32*100
X,Y = None, None
idx = 0
for x,y in iterate(shuffle=False):
print(x.shape, y.shape, x.dtype, y.dtype)
assert x.shape[0] == y.shape[0]
bs = x.shape[0]
if X is None:
X = Tensor.empty(sz, *x.shape[1:], device="disk:/tmp/imagenet_x", dtype=dtypes.uint8)
Y = Tensor.empty(sz, *y.shape[1:], device="disk:/tmp/imagenet_y", dtype=dtypes.int64)
print(X.shape, Y.shape)
X[idx:idx+bs].assign(x)
Y[idx:idx+bs].assign(y)
idx += bs
if idx >= sz: break