map config.architecture string to RNN Cell

Author: Ubuntu

examples/pytorch/FastCells/train_classifier.py

@@ -414,10 +414,10 @@ def create_model(model_config, input_size, num_keywords):
     wSparsity_list = [model_config.wSparsity, model_config.wSparsity, model_config.wSparsity]
     uSparsity_list = [model_config.uSparsity, model_config.uSparsity, model_config.uSparsity]
     print(model_config.gate_nonlinearity, model_config.update_nonlinearity)
-    return ModelClass(input_size, model_config.num_layers, hidden_units_list,
-                        wRank_list, uRank_list, wSparsity_list, uSparsity_list,
-                        model_config.gate_nonlinearity, model_config.update_nonlinearity,
-                        num_keywords)
+    return ModelClass(model_config.architecture, input_size, model_config.num_layers,
+                      hidden_units_list, wRank_list, uRank_list, wSparsity_list,
+                      uSparsity_list, model_config.gate_nonlinearity, 
+                      model_config.update_nonlinearity, num_keywords)
 
 def save_json(obj, filename):
     with open(filename, "w") as f:

pytorch/edgeml_pytorch/graph/rnn.py

@@ -1100,7 +1100,9 @@ def forward(self, input, hiddenState=None, cellState=None):
 class FastGRNNCUDA(nn.Module):
     """Unrolled implementation of the FastGRNNCUDACell"""
     def __init__(self, input_size, hidden_size, gate_nonlinearity="sigmoid",
-                 update_nonlinearity="tanh", wRank=None, uRank=None, zetaInit=1.0, nuInit=-4.0, name="FastGRNNCUDACell"):
+                 update_nonlinearity="tanh", wRank=None, uRank=None, 
+                 wSparsity=1.0, uSparsity=1.0, zetaInit=1.0, nuInit=-4.0,
+                 name="FastGRNNCUDACell"):
         super(FastGRNNCUDA, self).__init__()
         if utils.findCUDA() is None:
             raise Exception('FastGRNNCUDA is supported only on GPU devices.')
@@ -1309,4 +1311,4 @@ def forward(ctx, input, bias_gate, bias_update, zeta, nu, old_h, w, u, w1, w2, u
     def backward(ctx, grad_h):
         outputs = fastgrnn_cuda.backward_unroll(
             grad_h.contiguous(), *ctx.saved_variables, ctx.gate_non_linearity)
-        return tuple(outputs + [None])
+        return tuple(outputs + [None])

pytorch/edgeml_pytorch/trainer/fastmodel.py

@@ -14,10 +14,10 @@ class RNNClassifierModel(inheritance_class):
            RNN-based classifier
         """
 
-        def __init__(self, input_dim, num_layers, hidden_units_list,
+        def __init__(self, rnn_name, input_dim, num_layers, hidden_units_list,
                      wRank_list, uRank_list, wSparsity_list, uSparsity_list,
-                     gate_nonlinearity, update_nonlinearity,
-                     num_classes=None, linear=True, batch_first=False, apply_softmax=True):
+                     gate_nonlinearity, update_nonlinearity, num_classes=None,
+                     linear=True, batch_first=False, apply_softmax=True):
             """
             Initialize the KeywordSpotter with the following parameters:
             input_dim - the size of the input audio frame in # samples.
@@ -45,8 +45,9 @@ def __init__(self, input_dim, num_layers, hidden_units_list,
 
             super(RNNClassifierModel, self).__init__()
 
+            RNN = getattr(getattr(getattr(__import__('edgeml_pytorch'), 'graph'), 'rnn'), rnn_name)
             self.rnn_list = nn.ModuleList([
-                FastGRNN(self.input_dim if l==0 else self.hidden_units_list[l-1], 
+                RNN(self.input_dim if l==0 else self.hidden_units_list[l-1], 
                             self.hidden_units_list[l], 
                             gate_nonlinearity=self.gate_nonlinearity,
                             update_nonlinearity=self.update_nonlinearity,
@@ -144,4 +145,4 @@ def forward(self, input):
             if self.apply_softmax:
                 model_output = F.log_softmax(model_output, dim=1)
             return model_output
-    return RNNClassifierModel
+    return RNNClassifierModel