Import get_parameters from tinygrad.nn (#559)

* get_parameter is in optim

* Update all imports for get_parameters

* Clean up

* use optim.get_paramters

examples/deep_deterministic_policy_gradient.py

@@ -2,9 +2,8 @@ from typing import Optional, Tuple
 from numpy.typing import NDArray
 
 from tinygrad.tensor import Tensor
-import tinygrad.nn.optim as optim
+from tinygrad.nn import optim
 from tinygrad.helpers import getenv
-from extra.utils import get_parameters
 
 import numpy as np
 import gym
@@ -153,10 +152,10 @@ class DeepDeterministicPolicyGradient:
     self.target_actor = Actor(self.num_actions, self.num_states, hidden_size)
     self.target_critic = Critic(self.num_actions + self.num_states, hidden_size)
 
-    actor_params = get_parameters(self.actor)
-    critic_params = get_parameters(self.critic)
-    target_actor_params = get_parameters(self.target_actor)
-    target_critic_params = get_parameters(self.target_critic)
+    actor_params = optim.get_parameters(self.actor)
+    critic_params = optim.get_parameters(self.critic)
+    target_actor_params = optim.get_parameters(self.target_actor)
+    target_critic_params = optim.get_parameters(self.target_critic)
 
     if DEVICE == "GPU":
       [x.gpu_() for x in actor_params + critic_params + target_actor_params + target_critic_params]
@@ -172,12 +171,12 @@ class DeepDeterministicPolicyGradient:
       tau = self.tau
 
     for param, target_param in zip(
-        get_parameters(self.actor), get_parameters(self.target_actor)
+        optim.get_parameters(self.actor), optim.get_parameters(self.target_actor)
     ):
       target_param.assign(param * tau + target_param * (1.0 - tau))
 
     for param, target_param in zip(
-        get_parameters(self.critic), get_parameters(self.target_critic)
+        optim.get_parameters(self.critic), optim.get_parameters(self.target_critic)
     ):
       target_param.assign(param * tau + target_param * (1.0 - tau))
 

examples/hlb_cifar10.py

@@ -11,7 +11,6 @@ from tinygrad.nn import optim
 from tinygrad.tensor import Tensor
 from tinygrad.helpers import getenv
 from tinygrad.ops import GlobalCounters
-from extra.utils import get_parameters
 
 num_classes = 10
 
@@ -87,10 +86,10 @@ def train_cifar():
   Xt, Yt = fetch_batch(X_test, Y_test, BS=BS)
   model = SpeedyResNet()
   if getenv("ADAM"):
-    optimizer = optim.Adam(get_parameters(model), lr=Tensor([0.001]).realize())
+    optimizer = optim.Adam(optim.get_parameters(model), lr=Tensor([0.001]).realize())
   else:
-    #optimizer = optim.SGD(get_parameters(model), lr=0.001)
-    optimizer = optim.SGD(get_parameters(model), lr=Tensor([0.003]).realize(), momentum=0.85, nesterov=True)
+    #optimizer = optim.SGD(optim.get_parameters(model), lr=0.001)
+    optimizer = optim.SGD(optim.get_parameters(model), lr=Tensor([0.003]).realize(), momentum=0.85, nesterov=True)
 
   # 97 steps in 2 seconds = 20ms / step
   # step is 1163.42 GOPS = 56 TFLOPS!!!, 41% of max 136

examples/mnist_gan.py

@@ -7,8 +7,7 @@ import torch
 from torchvision.utils import make_grid, save_image
 from tinygrad.tensor import Tensor
 from tinygrad.helpers import getenv
-import tinygrad.nn.optim as optim
-from extra.utils import get_parameters
+from tinygrad.nn import optim
 from datasets import fetch_mnist
 GPU = getenv("GPU")
 
@@ -52,8 +51,8 @@ if __name__ == "__main__":
   batch_size = 512
   k = 1
   epochs = 300
-  generator_params = get_parameters(generator)
-  discriminator_params = get_parameters(discriminator)
+  generator_params = optim.get_parameters(generator)
+  discriminator_params = optim.get_parameters(discriminator)
   gen_loss = []
   disc_loss = []
   output_folder = "outputs"

examples/serious_mnist.py

@@ -7,7 +7,6 @@ from tinygrad.nn import BatchNorm2d, optim
 from tinygrad.helpers import getenv
 from datasets import fetch_mnist
 from extra.augment import augment_img
-from extra.utils import get_parameters
 from extra.training import train, evaluate, sparse_categorical_crossentropy
 GPU = getenv("GPU")
 QUICK = getenv("QUICK")
@@ -58,7 +57,7 @@ class BigConvNet:
 
   def parameters(self):
     if DEBUG: #keeping this for a moment
-      pars = [par for par in get_parameters(self) if par.requires_grad]
+      pars = [par for par in optim.get_parameters(self) if par.requires_grad]
       no_pars = 0
       for par in pars:
         print(par.shape)
@@ -66,17 +65,17 @@ class BigConvNet:
       print('no of parameters', no_pars)
       return pars
     else:
-      return get_parameters(self)
+      return optim.get_parameters(self)
 
   def save(self, filename):
     with open(filename+'.npy', 'wb') as f:
-      for par in get_parameters(self):
+      for par in optim.get_parameters(self):
         #if par.requires_grad:
         np.save(f, par.cpu().data)
 
   def load(self, filename):
     with open(filename+'.npy', 'rb') as f:
-      for par in get_parameters(self):
+      for par in optim.get_parameters(self):
         #if par.requires_grad:
         try:
           par.cpu().data[:] = np.load(f)
@@ -123,7 +122,7 @@ if __name__ == "__main__":
       print('could not load weights "'+sys.argv[1]+'".')
 
   if GPU:
-    params = get_parameters(model)
+    params = optim.get_parameters(model)
     [x.gpu_() for x in params]
 
   for lr, epochs in zip(lrs, epochss):

examples/train_efficientnet.py

@@ -3,12 +3,11 @@ import time
 from multiprocessing import Process, Queue
 import numpy as np
 from tqdm import trange
-import tinygrad.nn.optim as optim
+from tinygrad.nn import optim
 from tinygrad.helpers import getenv
 from tinygrad.tensor import Tensor
 from datasets import fetch_cifar
 from datasets.imagenet import fetch_batch
-from extra.utils import get_parameters
 from models.efficientnet import EfficientNet
 
 class TinyConvNet:
@@ -39,7 +38,7 @@ if __name__ == "__main__":
   else:
     model = EfficientNet(getenv("NUM", 0), classes, has_se=False)
 
-  parameters = get_parameters(model)
+  parameters = optim.get_parameters(model)
   print("parameter count", len(parameters))
   optimizer = optim.Adam(parameters, lr=0.001)
 

examples/train_resnet.py

@@ -2,9 +2,8 @@
 import numpy as np
 from PIL import Image
 
-from tinygrad.nn.optim import Adam
+from tinygrad.nn.optim import Adam, get_parameters
 from tinygrad.helpers import getenv
-from extra.utils import get_parameters
 from extra.training import train, evaluate
 from models.resnet import ResNet
 from datasets import fetch_mnist

examples/transformer.py

@@ -2,8 +2,7 @@
 import numpy as np
 import random
 
-from tinygrad.nn.optim import Adam
-from extra.utils import get_parameters
+from tinygrad.nn.optim import Adam, get_parameters
 from extra.training import train, evaluate
 from models.transformer import Transformer
 

extra/training.py

@@ -1,6 +1,5 @@
 import numpy as np
 from tqdm import trange
-from extra.utils import get_parameters
 from tinygrad.tensor import Tensor, Device
 from tinygrad.helpers import getenv
 

extra/utils.py

@@ -1,7 +1,7 @@
-from tinygrad.tensor import Tensor
 import pickle
-from tqdm import tqdm
 import numpy as np
+from tqdm import tqdm
+from tinygrad.tensor import Tensor
 from tinygrad.helpers import prod, getenv
 
 def fetch(url):
@@ -26,8 +26,6 @@ def download_file(url, fp, skip_if_exists=False):
       progress_bar.update(f.write(chunk))
   os.rename(fp+".tmp", fp)
 
-from tinygrad.nn.optim import get_parameters
-
 def my_unpickle(fb0):
   key_prelookup = {}
   class HackTensor:

test/graph_batchnorm.py

@@ -1,12 +1,11 @@
+import unittest
 from tinygrad.tensor import Tensor
 from tinygrad.nn import Conv2d, BatchNorm2d, optim
-from extra.utils import get_parameters  # TODO: move to optim
-import unittest
 
 def model_step(lm):
   Tensor.training = True
   x = Tensor.ones(8,12,128,256, requires_grad=False)
-  optimizer = optim.SGD(get_parameters(lm), lr=0.001)
+  optimizer = optim.SGD(optim.get_parameters(lm), lr=0.001)
   loss = lm.forward(x).sum()
   optimizer.zero_grad()
   loss.backward()

test/test_mnist.py

@@ -2,9 +2,8 @@
 import unittest
 import numpy as np
 from tinygrad.tensor import Tensor, Device
-import tinygrad.nn.optim as optim
+from tinygrad.nn import optim
 from extra.training import train, evaluate
-from extra.utils import get_parameters
 from datasets import fetch_mnist
 
 # load the mnist dataset
@@ -17,7 +16,7 @@ class TinyBobNet:
     self.l2 = Tensor.scaled_uniform(128, 10)
 
   def parameters(self):
-    return get_parameters(self)
+    return optim.get_parameters(self)
 
   def forward(self, x):
     return x.dot(self.l1).relu().dot(self.l2).logsoftmax()
@@ -34,7 +33,7 @@ class TinyConvNet:
     self.l1 = Tensor.scaled_uniform(out_chan*5*5, 10)
 
   def parameters(self):
-    return get_parameters(self)
+    return optim.get_parameters(self)
 
   def forward(self, x):
     x = x.reshape(shape=(-1, 1, 28, 28)) # hacks

test/test_optim.py

@@ -2,8 +2,7 @@ import numpy as np
 import torch
 import unittest
 from tinygrad.tensor import Tensor
-from tinygrad.nn.optim import Adam, SGD, RMSprop
-from extra.utils import get_parameters
+from tinygrad.nn.optim import Adam, SGD, RMSprop, get_parameters
 
 x_init = np.random.randn(1,3).astype(np.float32)
 W_init = np.random.randn(3,3).astype(np.float32)

test/test_train.py

@@ -1,11 +1,10 @@
 import unittest
 import time
-import tinygrad.nn.optim as optim
 import numpy as np
+from tinygrad.nn import optim
 from tinygrad.tensor import Device
 from tinygrad.helpers import getenv
 from extra.training import train
-from extra.utils import get_parameters
 from models.efficientnet import EfficientNet
 from models.transformer import Transformer
 from models.vit import ViT
@@ -14,7 +13,7 @@ from models.resnet import ResNet18
 BS = getenv("BS", 2)
 
 def train_one_step(model,X,Y):
-  params = get_parameters(model)
+  params = optim.get_parameters(model)
   pcount = 0
   for p in params:
     pcount += np.prod(p.shape)