Lambda layer improvements

Author: fchollet

keras/layers/core.py

@@ -12,9 +12,7 @@
 from .. import activations, initializations, regularizers, constraints
 from ..regularizers import ActivityRegularizer
 
-import marshal
-import types
-import sys
+import inspect
 
 
 class Layer(object):
@@ -1501,25 +1499,20 @@ class Lambda(Layer):
             Takes one argument: the output of previous layer
         output_shape: Expected output shape from function.
             Could be a tuple or a function of the shape of the input
+        arguments: optional dictionary of keyword arguments to be passed
+            to the function.
     '''
-    def __init__(self, function, output_shape=None, **kwargs):
+    def __init__(self, function, output_shape=None, arguments={}, **kwargs):
         super(Lambda, self).__init__(**kwargs)
-        py3 = sys.version_info[0] == 3
-        if py3:
-            self.function = marshal.dumps(function.__code__)
-        else:
-            assert hasattr(function, 'func_code'), ('The Lambda layer "function"'
-                                                    ' argument must be a Python function.')
-            self.function = marshal.dumps(function.func_code)
+        self.function = function
+        self.arguments = arguments
         if output_shape is None:
             self._output_shape = None
         elif type(output_shape) in {tuple, list}:
             self._output_shape = tuple(output_shape)
         else:
-            if py3:
-                self._output_shape = marshal.dumps(output_shape.__code__)
-            else:
-                self._output_shape = marshal.dumps(output_shape.func_code)
+            assert hasattr(output_shape, '__call__'), 'In Lambda, `output_shape` must be a list, a tuple, or a function.'
+            self._output_shape = output_shape
         super(Lambda, self).__init__()
 
     @property
@@ -1528,26 +1521,34 @@ def output_shape(self):
             # if TensorFlow, we can infer the output shape directly:
             if K._BACKEND == 'tensorflow':
                 # we assume output shape is not dependent on train/test mode
-                x = self.get_input()
+                x = self.get_output()
                 return K.int_shape(x)
             # otherwise, we default to the input shape
             return self.input_shape
-        elif type(self._output_shape) == tuple:
+        elif type(self._output_shape) in {tuple, list}:
             nb_samples = self.input_shape[0] if self.input_shape else None
-            return (nb_samples, ) + self._output_shape
+            return (nb_samples,) + tuple(self._output_shape)
         else:
-            output_shape_func = marshal.loads(self._output_shape)
-            output_shape_func = types.FunctionType(output_shape_func, globals())
-            shape = output_shape_func(self.input_shape)
+            shape = self._output_shape(self.input_shape)
             if type(shape) not in {list, tuple}:
                 raise Exception('output_shape function must return a tuple')
             return tuple(shape)
 
     def get_output(self, train=False):
         X = self.get_input(train)
-        func = marshal.loads(self.function)
-        func = types.FunctionType(func, globals())
-        return func(X)
+        arguments = self.arguments
+        arg_spec = inspect.getargspec(self.function)
+        if 'train' in arg_spec.args:
+            arguments['train'] = train
+        return self.function(X, **arguments)
+
+    def get_config(self):
+        # note: not serializable at the moment.
+        config = {'function': self.function,
+                  'output_shape': self._output_shape,
+                  'arguments': self.arguments}
+        base_config = super(Lambda, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
 
 
 class MaskedLambda(MaskedLayer, Lambda):
@@ -1567,8 +1568,10 @@ class LambdaMerge(Lambda):
             list of outputs from input layers
         output_shape - Expected output shape from function.
             Could be a tuple or a function of list of input shapes
+        arguments: optional dictionary of keyword arguments to be passed
+            to the function.
     '''
-    def __init__(self, layers, function, output_shape=None):
+    def __init__(self, layers, function, output_shape=None, arguments={}):
         if len(layers) < 2:
             raise Exception('Please specify two or more input layers '
                             '(or containers) to merge.')
@@ -1577,6 +1580,7 @@ def __init__(self, layers, function, output_shape=None):
         self.regularizers = []
         self.constraints = []
         self.updates = []
+        self.arguments = arguments
         for l in self.layers:
             params, regs, consts, updates = l.get_params()
             self.regularizers += regs
@@ -1586,45 +1590,39 @@ def __init__(self, layers, function, output_shape=None):
                 if p not in self.trainable_weights:
                     self.trainable_weights.append(p)
                     self.constraints.append(c)
-        py3 = sys.version_info[0] == 3
-        if py3:
-            self.function = marshal.dumps(function.__code__)
-        else:
-            self.function = marshal.dumps(function.func_code)
+        self.function = function
         if output_shape is None:
             self._output_shape = None
         elif type(output_shape) in {tuple, list}:
             self._output_shape = tuple(output_shape)
         else:
-            if py3:
-                self._output_shape = marshal.dumps(output_shape.__code__)
-            else:
-                self._output_shape = marshal.dumps(output_shape.func_code)
+            assert hasattr(output_shape, '__call__'), 'In LambdaMerge, `output_shape` must be a list, a tuple, or a function.'
+            self._output_shape = output_shape
         super(Lambda, self).__init__()
 
     @property
     def output_shape(self):
         input_shapes = [layer.output_shape for layer in self.layers]
         if self._output_shape is None:
             return input_shapes[0]
-        elif type(self._output_shape) == tuple:
-            return (input_shapes[0][0], ) + self._output_shape
+        elif type(self._output_shape) in {tuple, list}:
+            return (input_shapes[0][0],) + self._output_shape
         else:
-            output_shape_func = marshal.loads(self._output_shape)
-            output_shape_func = types.FunctionType(output_shape_func, globals())
-            shape = output_shape_func(input_shapes)
+            shape = self._output_shape(input_shapes)
             if type(shape) not in {list, tuple}:
-                raise Exception('output_shape function must return a tuple.')
+                raise Exception('In LambdaMerge, the `output_shape` function must return a tuple.')
             return tuple(shape)
 
     def get_params(self):
         return self.trainable_weights, self.regularizers, self.constraints, self.updates
 
     def get_output(self, train=False):
-        func = marshal.loads(self.function)
-        func = types.FunctionType(func, globals())
         inputs = [layer.get_output(train) for layer in self.layers]
-        return func(inputs)
+        arguments = self.arguments
+        arg_spec = inspect.getargspec(self.function)
+        if 'train' in arg_spec.args:
+            arguments['train'] = train
+        return self.function(inputs, **arguments)
 
     def get_input(self, train=False):
         res = []
@@ -1660,10 +1658,12 @@ def set_weights(self, weights):
             weights = weights[nb_param:]
 
     def get_config(self):
+        # note: not serializable at the moment.
         config = {'name': self.__class__.__name__,
                   'layers': [l.get_config() for l in self.layers],
                   'function': self.function,
-                  'output_shape': self._output_shape}
+                  'output_shape': self._output_shape,
+                  'arguments': self.arguments}
         base_config = super(LambdaMerge, self).get_config()
         return dict(list(base_config.items()) + list(config.items()))