add imagenet training support
ADD
datasets/imagenet.py
+import os
+import json
+import numpy as np
+from PIL import Image
+import torchvision.transforms as transforms
+
+BASEDIR = "/home/batman/imagenet"
+train_files = open(os.path.join(BASEDIR, "train_files")).read().split("\n")
+val_files = open(os.path.join(BASEDIR, "val_files")).read().split("\n")
+ci = json.load(open(os.path.join(BASEDIR, "imagenet_class_index.json")))
+cir = {v[0]: int(k) for k,v in ci.items()}
+
+rrc = transforms.RandomResizedCrop(224)
+def image_load(fn):
+ img = Image.open(fn).convert('RGB')
+ ret = np.array(rrc(img))
+ return ret
+
+def fetch_batch(bs, val=False):
+ files = val_files if val else train_files
+ samp = np.random.randint(0, len(files), size=(bs))
+ files = [files[i] for i in samp]
+ X = [image_load(os.path.join(BASEDIR, "val" if val else "train", x)) for x in files]
+ Y = [cir[x.split("/")[0]] for x in files]
+ return np.transpose(np.array(X), (0,3,1,2)), np.array(Y)
+
+if __name__ == "__main__":
+ X,Y = fetch_batch(64)
+ print(X.shape, Y)
+

MODIFY
examples/train_efficientnet.py
examples/train_efficientnet.py
@@ -28,6 +28,11 @@ class TinyConvNet:
return x.dot(self.l1).logsoftmax()
if __name__ == "__main__":
+ IMAGENET = os.getenv("IMAGENET") is not None
+ if IMAGENET:
+ from datasets.imagenet import fetch_batch
+ classes = 1000
+ else:
X_train, Y_train = fetch_cifar()
classes = 10
@@ -50,15 +55,18 @@ if __name__ == "__main__":
Tensor.training = True
for i in (t := trange(steps)):
+ if IMAGENET:
+ X,Y = fetch_batch(BS)
+ X = X.astype(np.float32)
+ else:
samp = np.random.randint(0, X_train.shape[0], size=(BS))
-
- img = X_train[samp].astype(np.float32)
+ X = X_train[samp].astype(np.float32)
+ Y = Y_train[samp]
st = time.time()
- out = model.forward(Tensor(img, requires_grad=False))
+ out = model.forward(Tensor(X, requires_grad=False))
fp_time = (time.time()-st)*1000.0
- Y = Y_train[samp]
y = np.zeros((BS,classes), np.float32)
y[range(y.shape[0]),Y] = -classes
y = Tensor(y, requires_grad=False)

add imagenet training support