Merge pull request #44 from torokati44/master

Allow using a backend other than TensorFlow

Author: minimaxir

textgenrnn/AttentionWeightedAverage.py

@@ -1,7 +1,6 @@
 from keras.engine import InputSpec, Layer
 from keras import backend as K
 from keras import initializers
-import tensorflow as tf
 
 
 class AttentionWeightedAverage(Layer):

textgenrnn/model.py

@@ -1,6 +1,6 @@
 from keras.optimizers import RMSprop
 from keras.layers import Input, Embedding, Dense, LSTM, Bidirectional
-from keras.layers import CuDNNLSTM, concatenate, Reshape, SpatialDropout1D
+from keras.layers import concatenate, Reshape, SpatialDropout1D
 from keras.models import Model
 from keras import backend as K
 from .AttentionWeightedAverage import AttentionWeightedAverage
@@ -68,8 +68,9 @@ def textgenrnn_model(num_classes, cfg, context_size=None,
 
 
 def new_rnn(cfg, layer_num):
-    has_gpu = len(K.tensorflow_backend._get_available_gpus()) > 0
-    if has_gpu:
+    use_cudnnlstm = K.backend() == 'tensorflow' and len(K.tensorflow_backend._get_available_gpus()) > 0
+    if use_cudnnlstm:
+        from keras.layers import CuDNNLSTM
         if cfg['rnn_bidirectional']:
             return Bidirectional(CuDNNLSTM(cfg['rnn_size'],
                                            return_sequences=True),

textgenrnn/utils.py

@@ -67,8 +67,8 @@ def textgenrnn_generate(model, vocab,
     if not isinstance(temperature, list):
         temperature = [temperature]
 
-    if model_input_count(model) > 1:
-        model = Model(inputs=model.input[0], outputs=model.output[1])
+    if len(model.inputs) > 1:
+        model = Model(inputs=model.inputs[0], outputs=model.outputs[1])
 
     while next_char != meta_token and len(text) < max_gen_length:
         encoded_text = textgenrnn_encode_sequence(text[-maxlen:],
@@ -166,13 +166,6 @@ def textgenrnn_encode_cat(chars, vocab):
     return a
 
 
-def model_input_count(model):
-    if isinstance(model.input, list):
-        return len(model.input)
-    else:   # is a Tensor
-        return model.input.shape[0]
-
-
 class generate_after_epoch(Callback):
     def __init__(self, textgenrnn, gen_epochs, max_gen_length):
         self.textgenrnn = textgenrnn
@@ -192,7 +185,7 @@ def __init__(self, weights_name, num_epochs, save_epochs):
         self.save_epochs = save_epochs
 
     def on_epoch_end(self, epoch, logs={}):
-        if model_input_count(self.model) > 1:
+        if len(self.model.inputs) > 1:
             self.model = Model(inputs=self.model.input[0],
                                outputs=self.model.output[1])
         if self.save_epochs > 0 and (epoch+1) % self.save_epochs == 0 and self.num_epochs != (epoch+1):