def test_integration(): from ape.examples.basic_computation import inputs, outputs, input_shapes from ape.examples.basic_computation import a, b, c, d, e from ape.examples.basic_network import machines, A, B from ape import timings comm_dict = load_dict("ape/test/integration_test_comm_dict.dat") comp_dict = load_dict("ape/test/integration_test_comp_dict.dat") rootdir = '_test/' os.system('mkdir -p %s' % rootdir) sanitize(inputs, outputs) known_shapes = shape_of_variables(inputs, outputs, input_shapes) comptime = timings.make_runtime_function(comp_dict) commtime = timings.make_commtime_function(comm_dict, known_shapes) assert isinstance(commtime(a, A, B), (int, float)) assert commtime(a, A, B) == commtime(a, B, A) elemwise = e.owner dot = d.owner assert comptime(elemwise, A) == 1 assert comptime(elemwise, B) == 100 assert comptime(dot, A) == 100 assert comptime(dot, B) == 1 graphs, scheds, rankfile, make = distribute(inputs, outputs, input_shapes, machines, commtime, comptime, 50) assert make == 18 # B-dot:1 + B-d-A:16 + A-+:1 # graphs == "{'A': ([b, a], [e]), 'B': ([a], [d])}" ais, [ao] = graphs[A] [bi], [bo] = graphs[B] assert set(map(str, ais)) == set("ab") assert ao.name == e.name assert bi.name == a.name assert bo.name == d.name assert rankfile[A] != rankfile[B] assert str(scheds['B'][0]) == str(dot) assert map(str, scheds['A']) == map(str, (c.owner, e.owner))
def test_integration(): from ape.examples.basic_computation import inputs, outputs, input_shapes from ape.examples.basic_computation import a,b,c,d,e from ape.examples.basic_network import machines, A, B from ape import timings comm_dict = load_dict("ape/test/integration_test_comm_dict.dat") comp_dict = load_dict("ape/test/integration_test_comp_dict.dat") rootdir = '_test/' os.system('mkdir -p %s'%rootdir) sanitize(inputs, outputs) known_shapes = shape_of_variables(inputs, outputs, input_shapes) comptime = timings.make_runtime_function(comp_dict) commtime = timings.make_commtime_function(comm_dict, known_shapes) assert isinstance(commtime(a, A, B), (int, float)) assert commtime(a, A, B) == commtime(a, B, A) elemwise = e.owner dot = d.owner assert comptime(elemwise, A) == 1 assert comptime(elemwise, B) == 100 assert comptime(dot, A) == 100 assert comptime(dot, B) == 1 graphs, scheds, rankfile, make = distribute(inputs, outputs, input_shapes, machines, commtime, comptime, 50) assert make == 18 # B-dot:1 + B-d-A:16 + A-+:1 # graphs == "{'A': ([b, a], [e]), 'B': ([a], [d])}" ais, [ao] = graphs[A] [bi], [bo] = graphs[B] assert set(map(str, ais)) == set("ab") assert ao.name == e.name assert bi.name == a.name assert bo.name == d.name assert rankfile[A] != rankfile[B] assert str(scheds['B'][0]) == str(dot) assert map(str, scheds['A']) == map(str, (c.owner, e.owner))
sanitize(inputs, outputs) # do timings if necessary recompute = False if recompute: comps = timings.comptime_dict(inputs, outputs, input_shapes, 5, machines, machine_groups) comms = timings.commtime_dict(network) save_dict(rootdir+'comps.dat', comps) save_dict(rootdir+'comms.dat', comms) else: comps = load_dict(rootdir+'comps.dat') comms = load_dict(rootdir+'comms.dat') known_shapes = shape_of_variables(inputs, outputs, input_shapes) comptime = timings.make_runtime_function(comps) commtime = timings.make_commtime_function(comms, known_shapes) # Break up graph graphs, scheds, rankfile, make = distribute(inputs, outputs, input_shapes, machines, commtime, comptime) # Write to disk write(graphs, scheds, rankfile, rootdir, known_shapes) # Print out fgraphs as pdfs fgraphs = {m: theano.FunctionGraph(*theano.gof.graph.clone(i, o)) for m, (i, o) in graphs.items()} for m, g in fgraphs.items(): theano.printing.pydotprint(g, outfile="%s%s.pdf"%(rootdir,m), format="pdf")
sanitize(inputs, outputs) # do timings if necessary recompute = False if recompute: comps = timings.comptime_dict(inputs, outputs, input_shapes, 5, machines, machine_groups) comms = timings.commtime_dict(network) save_dict(rootdir + 'comps.dat', comps) save_dict(rootdir + 'comms.dat', comms) else: comps = load_dict(rootdir + 'comps.dat') comms = load_dict(rootdir + 'comms.dat') known_shapes = shape_of_variables(inputs, outputs, input_shapes) comptime = timings.make_runtime_function(comps) commtime = timings.make_commtime_function(comms, known_shapes) # Break up graph graphs, scheds, rankfile, make = distribute(inputs, outputs, input_shapes, machines, commtime, comptime) # Write to disk write(graphs, scheds, rankfile, rootdir, known_shapes) # Print out fgraphs as pdfs fgraphs = { m: theano.FunctionGraph(*theano.gof.graph.clone(i, o)) for m, (i, o) in graphs.items() } for m, g in fgraphs.items():