Merge pull request #48 from nouiz/fix_pydotprint

theano/printing.py

@@ -382,7 +382,8 @@ pprint.assign(lambda pstate, r: hasattr(pstate, 'target') and pstate.target is n
 pp = pprint
 
 
-
+# colors not used: orange, amber#FFBF00, purple, pink,
+# used by default: green, blue, grey, red
 default_colorCodes = {'GpuFromHost' : 'red',
               'HostFromGpu' : 'red',
               'Scan'  : 'yellow',
@@ -391,13 +392,12 @@ default_colorCodes = {'GpuFromHost' : 'red',
               'Elemwise': '#FFAABB',
               'Subtensor': '#FFAAFF'}
 
-
-
 def pydotprint(fct, outfile=None,
                compact=True, format='png', with_ids=False,
                high_contrast=True, cond_highlight=None, colorCodes=None,
                max_label_size=50, scan_graphs=False,
-               var_with_name_simple=False
+               var_with_name_simple=False,
+               print_output_file=True
                ):
     """
     print to a file in png format the graph of op of a compile theano fct.
@@ -428,17 +428,21 @@ def pydotprint(fct, outfile=None,
                 we will print only the variable name.
                 Otherwise, we concatenate the type to the var name.
 
-    In the graph, box are an Apply Node(the execution of an op) and ellipse are variable.
-    If variable have name they are used as the text(if multiple var have the same name, they will be merged in the graph).
-    Otherwise, if the variable is constant, we print the value and finaly we print the type + an uniq number to don't have multiple var merged.
-    We print the op of the apply in the Apply box with a number that represent the toposort order of application of those Apply.
-    If an Apply have more then 1 input, print add a label to the edge that in the index of the inputs.
-
-    green ellipses are inputs to the graph
-    blue ellipses are outputs of the graph
-    grey ellipses are var generated by the graph that are not output and are not used.
-    red ellipses are transfer to/from the gpu.
-        op with those name GpuFromHost, HostFromGpu
+    In the graph, ellipses are Apply Nodes (the execution of an op)
+    and boxes are variables.  If variables have names they are used as
+    text (if multiple vars have the same name, they will be merged in
+    the graph).  Otherwise, if the variable is constant, we print its
+    value and finally we print the type + a unique number to prevent
+    multiple vars from being merged.  We print the op of the apply in
+    the Apply box with a number that represents the toposort order of
+    application of those Apply.  If an Apply has more than 1 input, we
+    label each edge between an input and the Apply node with the
+    input's index.
+
+    green boxes are inputs to the graph
+    blue boxes are outputs of the graph
+    grey boxes are vars generated by the graph that are not outputs and are not used
+    red ellipses are transfers from/to the gpu (ops with names GpuFromHost, HostFromGpu)
     """
     if colorCodes is None:
         colorCodes = default_colorCodes
@@ -450,7 +454,7 @@ def pydotprint(fct, outfile=None,
     if isinstance(fct, Function):
         mode = fct.maker.mode
         fct_env  = fct.maker.env
-        if not isinstance(mode,ProfileMode) or not mode.fct_call.has_key(fct):
+        if not isinstance(mode,ProfileMode) or not mode.profile_stats.has_key(fct):
             mode=None
     elif isinstance(fct, gof.Env):
         mode = None
@@ -542,14 +546,14 @@ def pydotprint(fct, outfile=None,
             return apply_name_cache[node]
         prof_str=''
         if mode:
-            time = mode.apply_time.get((topo.index(node),node),0)
+            time = mode.profile_stats[fct].apply_time.get(node,0)
             #second, % total time in profiler, %fct time in profiler
-            if mode.local_time[0]==0:
+            if mode.local_time==0:
                 pt=0
-            else: pt=time*100/mode.local_time[0]
-            if mode.fct_call[fct]==0:
+            else: pt=time*100/mode.local_time
+            if mode.profile_stats[fct].fct_callcount==0:
                 pf=0
-            else: pf = time*100/mode.fct_call_time[fct]
+            else: pf = time*100/mode.profile_stats[fct].fct_call_time
             prof_str='   (%.3fs,%.3f%%,%.3f%%)'%(time,pt,pf)
         applystr = str(node.op).replace(':','_')
         applystr += prof_str
@@ -659,7 +663,8 @@ def pydotprint(fct, outfile=None,
         outfile+='.'+format
     g.write(outfile, prog='dot', format=format)
 
-    print 'The output file is available at',outfile
+    if print_output_file:
+        print 'The output file is available at',outfile
     if scan_graphs:
         scan_ops = [(idx, x) for idx,x in enumerate(fct_env.toposort()) if isinstance(x.op, theano.scan_module.scan_op.Scan)]
         path, fn = os.path.split(outfile)

theano/tests/test_printing.py

@@ -29,7 +29,7 @@ def test_pydotprint_cond_highlight():
     theano.theano_logger.removeHandler(orig_handler)
     theano.theano_logger.addHandler(new_handler)
     try:
-        theano.printing.pydotprint(f, cond_highlight = True)
+        theano.printing.pydotprint(f, cond_highlight = True, print_output_file=False)
     finally:
         theano.theano_logger.addHandler(orig_handler)
         theano.theano_logger.removeHandler(new_handler)
@@ -37,6 +37,14 @@ def test_pydotprint_cond_highlight():
     assert s.getvalue() == 'pydotprint: cond_highlight is set but there is no IfElse node in the graph\n'
 
 
+def test_pydotprint_profile():
+    """Just check that pydotprint does not crash with ProfileMode."""
+
+    A = tensor.matrix()
+    f = theano.function([A], A+1, mode='ProfileMode')
+    theano.printing.pydotprint(f, print_output_file=False)
+
+
 def test_min_informative_str():
     """ evaluates a reference output to make sure the
         min_informative_str function works as intended """