Consistent testing (#137)

* Consistent GPU classes

Convert the existing GPU classes into one standard format.

Remove duplicated functions in `test_mnist` and create a TestMNISTGPU
class. This reduces line count and ensures consistency.

Use `@unittest.skipUnless(GPU, "Requires GPU")` instead of `if GPU:` to
skip GPU testing. This will ensure that skipped tests are displayed
accordingly in the pytest output.

* Optim Testing now supports GPU

* Tensor testing now supports GPU

jacobian and gradcheck auto skipped until GPU float64 support added.

* GPU support for custom constructor methods

* Remove GPU flag from Model constructors

It was requested that the `gpu` kwarg be removed from the model
constructor. GPU conversion is now handled in the train function.

This also required the conversion of Optimizer parameters as they are
constructed prior to execution of the `train` function and are dependant
on the model GPU state.

* Fix typo: float32->float64

* Clean `get_parameters` utility

Just a quick refactor w/ the new support for optimizers.

* Remove GPU kwarg from TinyNet

Remove `gpu` kwarg from tiny net to match test_mnist `train` function.

test/test_optim.py

@@ -1,20 +1,22 @@
 import numpy as np
 import torch
 import unittest
-from tinygrad.tensor import Tensor
+from tinygrad.tensor import Tensor, GPU
 from tinygrad.optim import Adam, SGD, RMSprop
+from tinygrad.utils import get_parameters 
 
 x_init = np.random.randn(1,3).astype(np.float32)
 W_init = np.random.randn(3,3).astype(np.float32)
 m_init = np.random.randn(1,3).astype(np.float32)
 
-def step_tinygrad(optim, kwargs={}):
+def step_tinygrad(optim, kwargs={}, gpu=False):
   net = TinyNet()
   optim = optim([net.x, net.W], **kwargs)
+  if gpu is True: [x.cuda_() for x in get_parameters([net, optim])]
   out = net.forward()
   out.backward()
   optim.step()
-  return net.x.data, net.W.data
+  return net.x.cpu().data, net.W.cpu().data
 
 def step_pytorch(optim, kwargs={}):
   net = TorchNet()
@@ -52,21 +54,29 @@ class TorchNet():
 
 
 class TestOptim(unittest.TestCase):
+  gpu = False
+
   def test_adam(self):
-    for x,y in zip(step_tinygrad(Adam),
+    for x,y in zip(step_tinygrad(Adam, gpu=self.gpu),
                    step_pytorch(torch.optim.Adam)):
       np.testing.assert_allclose(x, y, atol=1e-4)
 
   def test_sgd(self):
-    for x,y in zip(step_tinygrad(SGD, kwargs={'lr': 0.001}),
+    for x,y in zip(step_tinygrad(SGD, kwargs={'lr': 0.001}, gpu=self.gpu),
                    step_pytorch(torch.optim.SGD, kwargs={'lr': 0.001})):
       np.testing.assert_allclose(x, y, atol=1e-5)
 
   def test_rmsprop(self):
-    for x,y in zip(step_tinygrad(RMSprop, kwargs={'lr': 0.001, 'decay': 0.99}),
+    for x,y in zip(step_tinygrad(RMSprop, kwargs={'lr': 0.001, 'decay': 0.99}, gpu=self.gpu),
                    step_pytorch(torch.optim.RMSprop,
                                 kwargs={'lr': 0.001, 'alpha': 0.99})):
       np.testing.assert_allclose(x, y, atol=1e-5)
 
+
+@unittest.skipUnless(GPU, "Requires GPU")
+class TestOptimGPU(TestOptim):
+  gpu = True
+
+
 if __name__ == '__main__':
   unittest.main()