Real memory fix. Only enable LocalOptGroup profile for useless. But also keep…

theano/compile/mode.py

@@ -155,7 +155,7 @@ optdb.register('merge1', gof.MergeOptimizer(),
 # The opt should not do anything that need shape inference.
 # New nodes that don't have infer_shape need that the original node
 # also don't have infer_shape
-local_useless = gof.optdb.LocalGroupDB(apply_all_opts=True)
+local_useless = gof.optdb.LocalGroupDB(apply_all_opts=True, profile=True)
 optdb.register(
     'useless',
     gof.optdb.TopoDB(local_useless,

theano/gof/opt.py

@@ -1261,16 +1261,20 @@ class LocalOptGroup(LocalOptimizer):
                                   for opt in optimizers)
 
         self.apply_all_opts = kwargs.pop('apply_all_opts', False)
+        self.profile = kwargs.pop('profile', False)
         self.track_map = defaultdict(lambda: [])
         assert len(kwargs) == 0
+        if self.profile:
             self.time_opts = {}
+            self.time_nodes = {}
             self.process_count = {}
             self.applied_true = {}
             self.node_created = {}
 
         for o in self.opts:
-            self.process_count.setdefault(o, 0)
+            if self.profile:
                 self.time_opts.setdefault(o, 0)
+                self.time_nodes.setdefault(o, 0)
                 self.process_count.setdefault(o, 0)
                 self.applied_true.setdefault(o, 0)
                 self.node_created.setdefault(o, 0)
@@ -1303,11 +1307,13 @@ class LocalOptGroup(LocalOptimizer):
                 opt_start = time.time()
                 repl = opt.transform(node)
                 opt_finish = time.time()
+                if self.profile:
                     self.time_opts[opt] += opt_start - opt_finish
                     self.process_count[opt] += 1
                 if not repl:
                     continue
                 else:
+                    if self.profile:
                         self.node_created[opt] += len(graph.ops(fgraph.variables, repl))
                         self.applied_true[opt] += 1
                     if not multiple_opts or not repl[0].owner:
@@ -1322,11 +1328,19 @@ class LocalOptGroup(LocalOptimizer):
             return repl
 
         new_var = apply_mult_opts(node, node.fgraph, self.apply_all_opts)
+
+        node_finish = time.time()
+        if self.profile:
+            self.time_nodes[node] = node_finish - node_start
         return new_var
 
     @staticmethod
     def print_profile(stream, prof, level=0):
-        (time_opts, process_count, applied_true, node_created) = prof
+
+        if not self.profile:
+            return
+
+        (time_opts, time_nodes, process_count, applied_true, node_created) = prof
         blanc = ('    ' * int(level))
         print(blanc, "LocalOptGroup", file=stream)
         print(blanc, "---------------------", file=stream)

theano/gof/optdb.py

@@ -402,15 +402,18 @@ class LocalGroupDB(DB):
 
     """
 
-    def __init__(self, apply_all_opts=False):
+    def __init__(self, apply_all_opts=False, profile=False):
         super(LocalGroupDB, self).__init__()
         self.failure_callback = None
         self.apply_all_opts = apply_all_opts
+        self.profile = profile
 
     def query(self, *tags, **kwtags):
         # For the new `useless` optimizer
         opts = super(LocalGroupDB, self).query(*tags, **kwtags)
-        ret = opt.LocalOptGroup(*opts, apply_all_opts=self.apply_all_opts)
+        ret = opt.LocalOptGroup(*opts,
+                                apply_all_opts=self.apply_all_opts,
+                                profile=self.profile)
         return ret