Speed up softmax too

Update core.py Update core.py Update core.py Update core.py Update core.py Update theano_backend.py Update tensorflow_backend.py
jq · Feb 25, 2016 · 3aa807a · 3aa807a
1 parent 089fa11
commit 3aa807a
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 33 deletions.
diff --git a/keras/activations.py b/keras/activations.py
@@ -7,13 +7,9 @@ def softmax(x):
     if ndim == 2:
         return K.softmax(x)
     elif ndim == 3:
-        # apply softmax to each timestep
-        def step(x, states):
-            return K.softmax(x), []
-        last_output, outputs, states = K.rnn(step, x,
-                                             [],
-                                             mask=None)
-        return outputs
+        e = K.exp(x)
+        s = K.sum(e, axis=-1, keepdims=True)
+        return e / s
     else:
         raise Exception('Cannot apply softmax to a tensor that is not 2D or 3D. ' +
                         'Here, ndim=' + str(ndim))

diff --git a/keras/backend/tensorflow_backend.py b/keras/backend/tensorflow_backend.py
@@ -353,6 +353,8 @@ def spatial_2d_padding(x, padding=(1, 1), dim_ordering='th'):
                    [0, 0]]
     return tf.pad(x, pattern)
 
+def pack(x):
+    return tf.pack(x)
 
 # VALUE MANIPULATION
 

diff --git a/keras/backend/theano_backend.py b/keras/backend/theano_backend.py
@@ -415,6 +415,8 @@ def spatial_3d_padding(x, padding=(1, 1, 1), dim_ordering='th'):
         raise Exception('Invalid dim_ordering: ' + dim_ordering)
     return T.set_subtensor(output[indices], x)
 
+def pack(x):
+    return T.stack(*x)
 
 # VALUE MANIPULATION
 

diff --git a/keras/layers/core.py b/keras/layers/core.py
@@ -1121,34 +1121,19 @@ def output_shape(self):
 
     def get_output(self, train=False):
         X = self.get_input(train)
-        input_length = self.input_shape[1]
-        if input_length and K._BACKEND == 'theano':
+        if K._BACKEND == 'theano':
             import theano.tensor as T
-            #X: (nb_samples, timesteps, input_dim)
-            X = K.permute_dimensions(X, (1, 0, 2))
-            #X: (timesteps, nb_samples, input_dim)
-            W = [self.W] * input_length
-            W = T.stack(*W)
-            #W: (timesteps, input_dim, output_dim)
-            z = T.batched_tensordot(X, W, axes=[(2), (1)])
-            #z: (timesteps, nb_samples, output_dim)
-            z = K.permute_dimensions(z, (1, 0, 2))
-            #z: (nb_samples, timesteps, output_dim)
-            b = [self.b] * input_length
-            b = T.stack(*b)
-            #b: (timesteps, output_dim)
-            Y = self.activation(z + b)
+            z = T.tensordot(X, self.W, axes=[(2), (0)])
+            Y = z + self.b
+            Y = self.activation(Y)
+            return Y
+        elif K._BACKEND == 'tensorflow':
+            shape = self.input_shape
+            X = K.reshape(X, (-1, shape[-1]))
+            Y = K.dot(X, self.W) + self.b
+            Y = K.reshape(Y, (-1, shape[1], shape[2]))
+            Y = self.activation(Y)
             return Y
-
-        def step(x, states):
-            output = K.dot(x, self.W) + self.b
-            return output, []
-
-        last_output, outputs, states = K.rnn(step, X,
-                                             initial_states=[],
-                                             mask=None)
-        outputs = self.activation(outputs)
-        return outputs
 
     def get_config(self):
         config = {'name': self.__class__.__name__,