get parameters
MODIFY
examples/train_efficientnet.py
examples/train_efficientnet.py
@@ -3,10 +3,13 @@ import time
import numpy as np
from extra.efficientnet import EfficientNet
from tinygrad.tensor import Tensor
+from tinygrad.utils import get_parameters
if __name__ == "__main__":
Tensor.default_gpu = os.getenv("GPU") is not None
model = EfficientNet(int(os.getenv("NUM", "0")))
+ parameters = get_parameters(model)
+ print(len(parameters))
BS = 16
img = np.zeros((BS,3,224,224), dtype=np.float32)
@@ -26,6 +29,10 @@ if __name__ == "__main__":
y = Tensor(y)
loss = out.logsoftmax().mul(y).mean()
+ # zero grad
+ for p in parameters:
+ p.grad = None
+
st = time.time()
loss.backward()
et = time.time()

MODIFY
tinygrad/utils.py
tinygrad/utils.py
import numpy as np
+from tinygrad.tensor import Tensor
def layer_init_uniform(*x):
ret = np.random.uniform(-1., 1., size=x)/np.sqrt(np.prod(x))
@@ -18,3 +19,17 @@ def fetch(url):
os.rename(fp+".tmp", fp)
return dat
+def get_parameters(model):
+ parameters = []
+ if hasattr(model, '__dict__'):
+ for k,v in model.__dict__.items():
+ if isinstance(v, Tensor):
+ parameters.append(v)
+ elif isinstance(v, list):
+ for x in v:
+ parameters.extend(get_parameters(x))
+ elif hasattr(v, '__dict__'):
+ parameters.extend(get_parameters(v))
+ #print(k, type(v))
+ return parameters
+

get parameters