FastGRNNCUDA: Fixes (#136)

* fixes for installation and fastgrnncuda

* ensure input tensors are on device

* ensure tensors on device for fastgrnncudacell

* add batch_first support

* fix forward params

Author: MJ10

pytorch/edgeml_pytorch/graph/rnn.py

@@ -330,7 +330,8 @@ def __init__(self, input_size, hidden_size, gate_nonlinearity="sigmoid",
         self._zetaInit = zetaInit
         self._nuInit = nuInit
         self._name = name
-    
+        self.device = torch.device("cuda")
+
         if wRank is not None:
             self._num_W_matrices += 1
             self._num_weight_matrices[0] = self._num_W_matrices
@@ -340,29 +341,29 @@ def __init__(self, input_size, hidden_size, gate_nonlinearity="sigmoid",
         self._name = name
 
         if wRank is None:
-            self.W = nn.Parameter(0.1 * torch.randn([hidden_size, input_size]))
+            self.W = nn.Parameter(0.1 * torch.randn([hidden_size, input_size], self.device))
             self.W1 = torch.empty(0)
             self.W2 = torch.empty(0)
         else:
             self.W = torch.empty(0)
-            self.W1 = nn.Parameter(0.1 * torch.randn([wRank, input_size]))
-            self.W2 = nn.Parameter(0.1 * torch.randn([hidden_size, wRank]))
+            self.W1 = nn.Parameter(0.1 * torch.randn([wRank, input_size], self.device))
+            self.W2 = nn.Parameter(0.1 * torch.randn([hidden_size, wRank], self.device))
 
         if uRank is None:
-            self.U = nn.Parameter(0.1 * torch.randn([hidden_size, hidden_size]))
+            self.U = nn.Parameter(0.1 * torch.randn([hidden_size, hidden_size], self.device))
             self.U1 = torch.empty(0)
             self.U2 = torch.empty(0)
         else:
             self.U = torch.empty(0)
-            self.U1 = nn.Parameter(0.1 * torch.randn([uRank, hidden_size]))
-            self.U2 = nn.Parameter(0.1 * torch.randn([hidden_size, uRank]))
+            self.U1 = nn.Parameter(0.1 * torch.randn([uRank, hidden_size], self.device))
+            self.U2 = nn.Parameter(0.1 * torch.randn([hidden_size, uRank], self.device))
 
         self._gate_non_linearity = NON_LINEARITY[gate_nonlinearity]
 
-        self.bias_gate = nn.Parameter(torch.ones([1, hidden_size]))
-        self.bias_update = nn.Parameter(torch.ones([1, hidden_size]))
-        self.zeta = nn.Parameter(self._zetaInit * torch.ones([1, 1]))
-        self.nu = nn.Parameter(self._nuInit * torch.ones([1, 1]))
+        self.bias_gate = nn.Parameter(torch.ones([1, hidden_size], self.device))
+        self.bias_update = nn.Parameter(torch.ones([1, hidden_size], self.device))
+        self.zeta = nn.Parameter(self._zetaInit * torch.ones([1, 1], self.device))
+        self.nu = nn.Parameter(self._nuInit * torch.ones([1, 1], self.device))
 
     @property
     def name(self):
@@ -374,7 +375,11 @@ def cellType(self):
 
     def forward(self, input, state):
         # Calls the custom autograd function while invokes the CUDA implementation
-        return FastGRNNFunction.apply(input, self.bias_gate, self.bias_update, self.zeta, self.nu, h_state,
+        if not input.is_cuda:
+            input.to(self.device)
+        if not state.is_cuda:
+            state.to(self.device)
+        return FastGRNNFunction.apply(input, self.bias_gate, self.bias_update, self.zeta, self.nu, state,
             self.W, self.U, self.W1, self.W2, self.U1, self.U2, self._gate_non_linearity)
 
     def getVars(self):
@@ -1103,7 +1108,7 @@ class FastGRNNCUDA(nn.Module):
     def __init__(self, input_size, hidden_size, gate_nonlinearity="sigmoid",
                  update_nonlinearity="tanh", wRank=None, uRank=None, 
                  wSparsity=1.0, uSparsity=1.0, zetaInit=1.0, nuInit=-4.0,
-                 name="FastGRNNCUDACell"):
+                 batch_first=False, name="FastGRNNCUDA"):
         super(FastGRNNCUDA, self).__init__()
         if utils.findCUDA() is None:
             raise Exception('FastGRNNCUDA is supported only on GPU devices.')
@@ -1113,7 +1118,17 @@ def __init__(self, input_size, hidden_size, gate_nonlinearity="sigmoid",
         self._zetaInit = zetaInit
         self._nuInit = nuInit
         self._name = name
-    
+        self._num_W_matrices = 1
+        self._num_U_matrices = 1
+        self._num_biases = 2
+        self._num_weight_matrices = [self._num_W_matrices, self._num_U_matrices, self._num_biases]
+        self._wRank = wRank
+        self._uRank = uRank
+        self._wSparsity = wSparsity
+        self._uSparsity = uSparsity
+        self.oldmats = []
+        self.device = torch.device("cuda")
+        self.batch_first = batch_first
         if wRank is not None:
             self._num_W_matrices += 1
             self._num_weight_matrices[0] = self._num_W_matrices
@@ -1123,33 +1138,42 @@ def __init__(self, input_size, hidden_size, gate_nonlinearity="sigmoid",
         self._name = name
 
         if wRank is None:
-            self.W = nn.Parameter(0.1 * torch.randn([hidden_size, input_size]))
+            self.W = nn.Parameter(0.1 * torch.randn([hidden_size, input_size], device=self.device))
             self.W1 = torch.empty(0)
             self.W2 = torch.empty(0)
         else:
             self.W = torch.empty(0)
-            self.W1 = nn.Parameter(0.1 * torch.randn([wRank, input_size]))
-            self.W2 = nn.Parameter(0.1 * torch.randn([hidden_size, wRank]))
+            self.W1 = nn.Parameter(0.1 * torch.randn([wRank, input_size], device=self.device))
+            self.W2 = nn.Parameter(0.1 * torch.randn([hidden_size, wRank], device=self.device))
 
         if uRank is None:
-            self.U = nn.Parameter(0.1 * torch.randn([hidden_size, hidden_size]))
+            self.U = nn.Parameter(0.1 * torch.randn([hidden_size, hidden_size], device=self.device))
             self.U1 = torch.empty(0)
             self.U2 = torch.empty(0)
         else:
             self.U = torch.empty(0)
-            self.U1 = nn.Parameter(0.1 * torch.randn([uRank, hidden_size]))
-            self.U2 = nn.Parameter(0.1 * torch.randn([hidden_size, uRank]))
+            self.U1 = nn.Parameter(0.1 * torch.randn([uRank, hidden_size], device=self.device))
+            self.U2 = nn.Parameter(0.1 * torch.randn([hidden_size, uRank], device=self.device))
 
         self._gate_non_linearity = NON_LINEARITY[gate_nonlinearity]
 
-        self.bias_gate = nn.Parameter(torch.ones([1, hidden_size]))
-        self.bias_update = nn.Parameter(torch.ones([1, hidden_size]))
-        self.zeta = nn.Parameter(self._zetaInit * torch.ones([1, 1]))
-        self.nu = nn.Parameter(self._nuInit * torch.ones([1, 1]))
+        self.bias_gate = nn.Parameter(torch.ones([1, hidden_size], device=self.device))
+        self.bias_update = nn.Parameter(torch.ones([1, hidden_size], device=self.device))
+        self.zeta = nn.Parameter(self._zetaInit * torch.ones([1, 1], device=self.device))
+        self.nu = nn.Parameter(self._nuInit * torch.ones([1, 1], device=self.device))
 
-    def forward(self, input, h_state, cell_state=None):
+    def forward(self, input, hiddenState, cell_state=None):
         # input: [timesteps, batch, features, state_size]
-        return FastGRNNUnrollFunction.apply(input, self.bias_gate, self.bias_update, self.zeta, self.nu, h_state,
+        if self.batch_first:
+            input = input.transpose(0, 1)
+        if not input.is_cuda:
+            input = input.to(self.device)
+        if hiddenState is None:
+            hiddenState = torch.zeros(
+                [input.shape[1], self.hidden_size]).to(self.device)
+        if not hiddenState.is_cuda:
+            hiddenState = hiddenState.to(self.device)
+        return FastGRNNUnrollFunction.apply(input, self.bias_gate, self.bias_update, self.zeta, self.nu, hiddenState,
             self.W, self.U, self.W1, self.W2, self.U1, self.U2, self._gate_non_linearity)
 
     def getVars(self):

pytorch/setup.py

@@ -22,7 +22,7 @@
     version='0.3.0',
     description='PyTorch code for ML algorithms for edge devices developed at Microsoft Research India.',
     author_email="edgeml@microsoft.com",
-    packages=['edgeml_pytorch'],
+    packages=['edgeml_pytorch', 'edgeml_pytorch.trainer', 'edgeml_pytorch.graph'],
     license='MIT License',
     long_description=open('README.md').read(),
     url='https://github.com/Microsoft/EdgeML',