FastCell Example Fixes, Generalized trainer for both batch_first args

examples/pytorch/FastCells/fastcell_example.py

@@ -11,7 +11,7 @@
 
 def main():
     # change cuda:0 to cuda:gpuid for specific allocation
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    device = torch.device("cuda:3" if torch.cuda.is_available() else "cpu")
     # Fixing seeds for reproducibility
     torch.manual_seed(42)
     np.random.seed(42)
@@ -43,10 +43,16 @@ def main():
 
     (dataDimension, numClasses, Xtrain, Ytrain, Xtest, Ytest,
      mean, std) = helpermethods.preProcessData(dataDir)
-
     assert dataDimension % inputDims == 0, "Infeasible per step input, " + \
         "Timesteps have to be integer"
 
+    timeSteps = int(Xtest.shape[1] / inputDims)
+    Xtest = np.reshape(Xtest, (-1, timeSteps, inputDims))
+    Xtrain = Xtrain.reshape((-1, timeSteps, inputDims))
+    if not args.batch_first:
+        Xtest = np.swapaxes(Xtest, 0, 1)
+        Xtrain = np.swapaxes(Xtrain, 0, 1)
+
     currDir = helpermethods.createTimeStampDir(dataDir, cell)
 
     helpermethods.dumpCommand(sys.argv, currDir)

pytorch/edgeml_pytorch/trainer/fastTrainer.py

@@ -77,9 +77,13 @@ def computeLogits(self, input):
             logits = self.classifier(feats[-1, :])
         else:
             feats = self.RNN(input)
-            logits = self.classifier(feats[-1, :])
-
-        return logits, feats[:, -1]
+            if self.batch_first:
+                logits = self.classifier(feats[:, -1])
+                feats_n = feats[:,-1]
+            else:
+                logits = self.classifier(feats[-1,:])
+                feats_n = feats[-1,:]
+        return logits, feats_n
 
     def optimizer(self):
         '''
@@ -351,7 +355,13 @@ def train(self, batchSize, totalEpochs, Xtrain, Xtest, Ytrain, Ytest,
         '''
         fileName = str(self.FastObj.cellType) + 'Results_pytorch.txt'
         resultFile = open(os.path.join(dataDir, fileName), 'a+')
-        numIters = int(np.ceil(float(Xtrain.shape[0]) / float(batchSize)))
+        if self.batch_first:
+                self.timeSteps = Xtrain.shape[1]
+                self.numPoints = Xtrain.shape[0]
+        else:
+                self.timeSteps = Xtrain.shape[0]
+                self.numPoints = Xtrain.shape[1]
+        numIters = int(np.ceil(float(self.numPoints) / float(batchSize)))
         totalBatches = numIters * totalEpochs
 
         counter = 0
@@ -362,11 +372,6 @@ def train(self, batchSize, totalEpochs, Xtrain, Xtest, Ytrain, Ytest,
             ihtDone = 1
             maxTestAcc = -10000
         header = '*' * 20
-        self.timeSteps = int(Xtest.shape[1] / self.inputDims)
-        Xtest = Xtest.reshape((-1, self.timeSteps, self.inputDims))
-        Xtest = np.swapaxes(Xtest, 0, 1)
-        Xtrain = Xtrain.reshape((-1, self.timeSteps, self.inputDims))
-        Xtrain = np.swapaxes(Xtrain, 0, 1)
 
         for i in range(0, totalEpochs):
             print("\nEpoch Number: " + str(i), file=self.outFile)
@@ -376,7 +381,7 @@ def train(self, batchSize, totalEpochs, Xtrain, Xtest, Ytrain, Ytest,
                 for param_group in self.optimizer.param_groups:
                     param_group['lr'] = self.learningRate
 
-            shuffled = list(range(Xtrain.shape[1]))
+            shuffled = list(range(self.numPoints))
             np.random.shuffle(shuffled)
             trainAcc = 0.0
             trainLoss = 0.0
@@ -389,9 +394,12 @@ def train(self, batchSize, totalEpochs, Xtrain, Xtest, Ytrain, Ytest,
                           (header, msg, header), file=self.outFile)
 
                 k = shuffled[j * batchSize:(j + 1) * batchSize]
-                batchX = Xtrain[:, k, :]
+                if self.batch_first:
+                        batchX = Xtrain[k, :, :]
+                else:
+                        batchX = Xtrain[:, k, :]
+
                 batchY = Ytrain[k]
-
                 self.optimizer.zero_grad()
                 logits, _ = self.computeLogits(batchX.to(self.device))
                 batchLoss = self.loss(logits, batchY.to(self.device))