Change the reset method of viewed_by

theano/compile/profiling.py

@@ -819,9 +819,13 @@ class ProfileStats(object):
                     # We don't want a shallow copy, but we don't want
                     # a deep copy. So this do a "middle" copy, where
                     # we copy the dict and the list, but not the var
-                    viewed_by_temp = {}
-                    for k, v in viewed_by.iteritems():
-                        viewed_by_temp[k] = list(v)
+                    # viewed_by_temp = {}
+                    # for k, v in viewed_by.iteritems():
+                    #     viewed_by_temp[k] = list(v)
+
+                    change_track_add = defaultdict(lambda: [])
+                    change_track_remove = defaultdict(lambda: [])
+                    # use to track viewed_by changes
 
                     for var in node.outputs:
                         compute_map[var][0] = 1
@@ -856,7 +860,8 @@ class ProfileStats(object):
                             # We keep trac of view only again the original
                             origin = view_of_temp.get(ins, ins)
                             view_of_temp[out] = origin
-                            viewed_by_temp[origin].append(out)
+                            viewed_by[origin].append(out)
+                            change_track_add[origin].append(out)
                         else:
                             mem_created += var_mem[out]
                         idx += 1
@@ -867,18 +872,19 @@ class ProfileStats(object):
                     # Mimic the combination of Theano and Python gc.
                     for ins in node.inputs:
                         assert not (ins in view_of_temp and
-                                    viewed_by_temp[ins])
+                                    viewed_by[ins])
                         # We track of the original var, so this shouldn't happen
                         if (dependencies[ins] and
                             ins not in fgraph.outputs and
                             ins.owner and
                             all([compute_map[v][0] for v in dependencies[ins]])):
-                            if ins not in view_of_temp and not viewed_by_temp.get(ins, []):
+                            if ins not in view_of_temp and not viewed_by.get(ins, []):
                                 mem_freed += var_mem[ins]
                             elif ins in view_of_temp:
                                 origin = view_of_temp[ins]
-                                viewed_by_temp[origin].remove(ins)
-                                if (not viewed_by_temp[origin] and
+                                viewed_by[origin].remove(ins)
+                                change_track_remove[origin].append(ins)
+                                if (not viewed_by[origin] and
                                     origin not in fgraph.inputs and
                                     not isinstance(origin, theano.Constant)):
                                     mem_freed += var_mem[origin]
@@ -899,7 +905,7 @@ class ProfileStats(object):
                         if max_mem_count < mem_bound:
                             mem_bound = max_mem_count
                     else:
-                        min_memory_generator(new_exec_nodes, viewed_by_temp, view_of_temp)
+                        min_memory_generator(new_exec_nodes, viewed_by, view_of_temp)
 
                     # Reset track variables
                     mem_count -= mem_created
@@ -908,6 +914,15 @@ class ProfileStats(object):
                     for var in node.outputs:
                         compute_map[var][0] = 0
 
+                    for k_remove, v_remove in change_track_remove.iteritems():
+                        for i in v_remove:
+                            viewed_by[k_remove].append(i)
+
+                    for k_add, v_add in change_track_add.iteritems():
+                        for i in v_add:
+                            viewed_by[k_add].remove(i)
+
+
             # two data structure used to mimic Python gc
             viewed_by = {}  # {var1: [vars that view var1]}
             # The len of the list is the value of python ref count. But we use a list, not just the ref count value.