support: no need to re-optmize when unpickling theano function. added a new…

theano/compile/function_module.py

@@ -703,6 +703,7 @@ class Function(object):
 # pickling/deepcopy support for Function
 
 def _pickle_Function(f):
+    print 'pickling Function..'
     #copy of the input storage list
     ins = list(f.input_storage)
     input_storage = []
@@ -736,11 +737,11 @@ def _pickle_Function(f):
                                     ' operation') %(str(d_i), str(d_j)))
                         else:
                             raise AliasedMemoryError(d_i, d_j)
-
     rval = (_constructor_Function, (f.maker, input_storage, inputs_data))
     return rval
 
 def _constructor_Function(maker, input_storage, inputs_data):
+    print 'unpickling Function...'
     if not theano.config.unpickle_function:
         return None
     f = maker.create(input_storage, trustme = True)
@@ -970,103 +971,9 @@ class FunctionMaker(object):
         else:
             raise TypeError("Unknown output type: %s (%s)", type(output), output)
         
-    def __init__(self, inputs, outputs,
-            mode=None, accept_inplace=False, function_builder=Function,
-            profile=None, on_unused_input=None):
-        """
-        :type inputs: a list of SymbolicInput instances
-
-        :type outputs: a list of SymbolicOutput instances
-                    outputs may also be a single Variable (not a list), in which
-                    case the functions produced by FunctionMaker will return
-                    their output value directly
-
-        :param mode: a Mode instance telling FunctionMaker how to optimize and link.  None
-            means to use the `config.mode`.
-
-        :param accept_inplace: True iff it is acceptable to have inplace operations
-            in the graph from the inputs to the outputs
-
-        :param on_unused_input: What to do if a variable in the 'inputs' list
-            is not used in the graph. Possible values are:
-                - 'raise': raise an error
-                - 'warn': log a warning
-                - 'ignore': do not do anything
-                - None: Use the value in the Theano flags on_unused_input
-        """
-        mode = theano.compile.mode.get_mode(mode)
-
-        # figure out which profile object to use (if any)
-        # to help with forward-porting ProfileMode,
-        # we allow ProfileMode to provide a ProfileStats object
-        # using this somewhat awkward mechanism.
-        mode_profile = getattr(mode, 'profile', None)
-        if (profile is not None and
-            profile is not False and
-            mode_profile is not None):
-            raise TypeError(
-                    'profile passed via both "mode" and "profile" arguments')
-        self.profile = profile = profile or mode_profile
-        if profile or theano.config.cxx:
-            # This is very important:
-            # 1) We preload the cache here to don't have its timming
-            #    included in optimization that compile function.
-            # 2) If other repo that import Theano have Theano ops defined,
-            #    we need to refresh the cache here. Otherwise, their is import
-            #    order problems.
-            #    When device=gpu, we compile during Theano import. This trigger
-            #    the loading of the cache. But unpickling the cache ask that the
-            #    other repos Ops are completly loaded, which isn't always the
-            #    case!
-            #    If a module isn't completly loaded and their unpickling fail,
-            #    it mean it is safe for this function compilation to skip them,
-            #    but not for futur compilation. So reloading the cache at each
-            #    compilation fix this problem.
-            # 3) This help propagate knowledge of newly compiled module to
-            #    concurrent process.
-            theano.gof.cc.get_module_cache().refresh()
-        # Handle the case where inputs and/or outputs is a single Variable (not in a list)
-        self.orig_outputs = outputs
-        unpack_single = False
-        return_none = False
-        if outputs is None:
-            return_none = True
-            outputs = []
-        if not isinstance(outputs, (list, tuple)):
-            unpack_single = True
-            outputs = [outputs]
-        if not isinstance(inputs, (list, tuple)):
-            inputs = [inputs]
-
-        # Wrap them in In or Out instances if needed.
-        #import pudb; pudb.set_trace()
-        inputs, outputs = map(self.wrap_in, inputs), map(self.wrap_out, outputs)
-        _inputs = gof.graph.inputs([o.variable for o in outputs] + [i.update
-            for i in inputs if getattr(i, 'update', False)])
-
-        # Check if some input variables are unused
-        self._check_unused_inputs(inputs, outputs, on_unused_input)
-
-        # Make a list of (SymbolicInput|SymblicInputKits, indices, [SymbolicInput,...]), one 
-        # tuple for each input. (See Function.indices for more details)
-        indices = [[input] + self.expand_in(input, _inputs) for input in inputs]
-
-        # make the fgraph (copies the graph, creates NEW INPUT AND OUTPUT VARIABLES)
-        fgraph, additional_outputs = std_fgraph(inputs, outputs, accept_inplace)
-        fgraph.profile = profile
-
-        self.fgraph = fgraph
-
-        # Fetch the optimizer and linker
-        optimizer, linker = mode.optimizer, copy.copy(mode.linker)
-
-        # optimize the fgraph
-        compute_test_value_orig = theano.config.compute_test_value
-        add_stack_trace_on_call = gof.Op.add_stack_trace_on_call
-        try:
-            theano.config.compute_test_value = theano.config.compute_test_value_opt
-            gof.Op.add_stack_trace_on_call = False
-
+    def retrieve_fgraph_from_opt_cache():
+        # This function is not finished
+        raise NotImplementedError('optimization cache is not finished! Should not be called.')
         from theano.gof.compilelock import get_lock, release_lock
         import os.path
         graph_db_file = os.path.join(theano.config.compiledir, 'optimized_graphs.pkl')
@@ -1203,8 +1110,6 @@ class FunctionMaker(object):
                         need_optimize = False
                         key = graph_old
                         break
-
-                # now optimize or not
                 if need_optimize:
                     # this is a brand new graph, optimize it, save it to graph_db
                     print 'optimizing the graph'
@@ -1235,15 +1140,119 @@ class FunctionMaker(object):
                 # release stuff
                 release_lock()
                 
-            #end of cache optimization
                 
-            #else containing the old code
+    def __init__(self, inputs, outputs,
+            mode=None, accept_inplace=False, function_builder=Function,
+            profile=None, on_unused_input=None, fgraph=None):
+        """
+        :type inputs: a list of SymbolicInput instances
+
+        :type outputs: a list of SymbolicOutput instances
+                    outputs may also be a single Variable (not a list), in which
+                    case the functions produced by FunctionMaker will return
+                    their output value directly
+
+        :param mode: a Mode instance telling FunctionMaker how to optimize and link.  None
+            means to use the `config.mode`.
+
+        :param accept_inplace: True iff it is acceptable to have inplace operations
+            in the graph from the inputs to the outputs
+
+        :param on_unused_input: What to do if a variable in the 'inputs' list
+            is not used in the graph. Possible values are:
+                - 'raise': raise an error
+                - 'warn': log a warning
+                - 'ignore': do not do anything
+                - None: Use the value in the Theano flags on_unused_input
+        """
+        mode = theano.compile.mode.get_mode(mode)
+
+        # figure out which profile object to use (if any)
+        # to help with forward-porting ProfileMode,
+        # we allow ProfileMode to provide a ProfileStats object
+        # using this somewhat awkward mechanism.
+        mode_profile = getattr(mode, 'profile', None)
+        if (profile is not None and
+            profile is not False and
+            mode_profile is not None):
+            raise TypeError(
+                    'profile passed via both "mode" and "profile" arguments')
+        self.profile = profile = profile or mode_profile
+        if profile or theano.config.cxx:
+            # This is very important:
+            # 1) We preload the cache here to don't have its timming
+            #    included in optimization that compile function.
+            # 2) If other repo that import Theano have Theano ops defined,
+            #    we need to refresh the cache here. Otherwise, their is import
+            #    order problems.
+            #    When device=gpu, we compile during Theano import. This trigger
+            #    the loading of the cache. But unpickling the cache ask that the
+            #    other repos Ops are completly loaded, which isn't always the
+            #    case!
+            #    If a module isn't completly loaded and their unpickling fail,
+            #    it mean it is safe for this function compilation to skip them,
+            #    but not for futur compilation. So reloading the cache at each
+            #    compilation fix this problem.
+            # 3) This help propagate knowledge of newly compiled module to
+            #    concurrent process.
+            theano.gof.cc.get_module_cache().refresh()
+        # Handle the case where inputs and/or outputs is a single Variable (not in a list)
+        self.orig_outputs = outputs
+        unpack_single = False
+        return_none = False
+        if outputs is None:
+            return_none = True
+            outputs = []
+        if not isinstance(outputs, (list, tuple)):
+            unpack_single = True
+            outputs = [outputs]
+        if not isinstance(inputs, (list, tuple)):
+            inputs = [inputs]
+
+        # Wrap them in In or Out instances if needed.
+        inputs, outputs = map(self.wrap_in, inputs), map(self.wrap_out, outputs)
+        _inputs = gof.graph.inputs([o.variable for o in outputs] + [i.update
+            for i in inputs if getattr(i, 'update', False)])
+
+        # Check if some input variables are unused
+        self._check_unused_inputs(inputs, outputs, on_unused_input)
+
+        # Make a list of (SymbolicInput|SymblicInputKits, indices, [SymbolicInput,...]), one 
+        # tuple for each input. (See Function.indices for more details)
+        indices = [[input] + self.expand_in(input, _inputs) for input in inputs]
+
+        if fgraph is None:
+            need_opt = True
+        else:
+            need_opt = False
+            
+        if fgraph is None:
+            # make the fgraph (copies the graph, creates NEW INPUT AND OUTPUT VARIABLES)
+            fgraph, additional_outputs = std_fgraph(inputs, outputs, accept_inplace)
+            fgraph.profile = profile
         else:
+            # fgraph is already an optimized one
+            _, additional_outputs = std_fgraph(inputs, outputs, accept_inplace)
+            pass
+        
+        self.fgraph = fgraph
+
+        # Fetch the optimizer and linker
+        optimizer, linker = mode.optimizer, copy.copy(mode.linker)
+
+        compute_test_value_orig = theano.config.compute_test_value
+        add_stack_trace_on_call = gof.Op.add_stack_trace_on_call
+        
+        if need_opt:
+            # optimize the fgraph
+            print 'fgraph is optimized'
+            try:
+                theano.config.compute_test_value = theano.config.compute_test_value_opt
+                gof.Op.add_stack_trace_on_call = False
                 start_optimizer = time.time()
                 optimizer_profile = optimizer(fgraph)
                 end_optimizer = time.time()
                 opt_time = end_optimizer - start_optimizer
-            #print 'opt took %s' % opt_time
                 if profile:
                     profile.optimizer_time += opt_time
                     if theano.config.profile_optimizer:
@@ -1256,6 +1265,12 @@ class FunctionMaker(object):
                 theano.config.compute_test_value = compute_test_value_orig
                 gof.Op.add_stack_trace_on_call = add_stack_trace_on_call
                 
+        else:
+            # fgraph is already optimized
+            print 'fgraph is not optimized'
+            theano.config.compute_test_value = compute_test_value_orig
+            gof.Op.add_stack_trace_on_call = add_stack_trace_on_call
+        
         # initialize the linker
         if not hasattr(linker, 'accept'):
             raise ValueError("'linker' parameter of FunctionMaker should be a Linker with an accept method " \
@@ -1415,9 +1430,11 @@ class FunctionMaker(object):
 
 
 def _pickle_FunctionMaker(self):
+    'picking FunctionMaker'
     kwargs = dict(
                 inputs=self.inputs,
                 outputs=self.orig_outputs,
+                fgraph=self.fgraph,
                 mode=self.mode,
                 accept_inplace=self.accept_inplace,
                 function_builder=self.function_builder,
@@ -1428,7 +1445,10 @@ def _pickle_FunctionMaker(self):
 
 
 def _constructor_FunctionMaker(kwargs):
+    print 'unpickling FunctionMaker...'
     if theano.config.unpickle_function:
+        if theano.config.reoptimize_unpickled_function:
+            del kwargs['fgraph']
         return FunctionMaker(**kwargs)
     else:
         return None

theano/configdefaults.py

@@ -118,6 +118,7 @@ AddConfigVar('print_active_device',
         BoolParam(True, allow_override=False),
         in_c_key=False)
 
+
 # Do not add FAST_RUN_NOGC to this list (nor any other ALL CAPS shortcut).
 # The way to get FAST_RUN_NOGC is with the flag 'linker=c|py_nogc'.
 # The old all capital letter way of working is deprecated as it is not
@@ -465,6 +466,12 @@ AddConfigVar('unpickle_function',
              BoolParam(True),
              in_c_key=False)
 
+AddConfigVar('reoptimize_unpickled_function',
+        "Re-optimize the graph when a theano function is unpickled from the disk.",
+        BoolParam(False, allow_override=False),
+        in_c_key=False)
+
+
 """Note to developers:
     Generally your exceptions should use an apply node's __str__
     method when exception_verbosity == 'low'. When exception_verbosity

theano/tests/test_pickle.py

@@ -8,7 +8,7 @@ floatX = 'float32'
 
 def test_pickle_unpickle():
     # Test if pick and unpickling a theano function with
-    # shared variables work
+    # shared variables should be pickled properly
     x1 = T.fmatrix('x1')
     x2 = T.fmatrix('x2')
     x3 = theano.shared(numpy.ones((10,10),dtype=floatX))
@@ -24,7 +24,7 @@ def test_pickle_unpickle():
     cPickle.dump(f, pkl_path, -1)
     pkl_path = open('thean_fn.pkl','r')
     f_ = cPickle.load(pkl_path)
-
+    import ipdb; ipdb.set_trace()
     in1 = numpy.ones((10, 10), dtype=floatX)
     in2 = numpy.ones((10, 10), dtype=floatX)