def test_create_plan_from_proto_correctly(self): from caffe2.python.net_builder import ops with Node('trainer'), Task(name='my_task', num_instances=2) as task: with ops.task_init(): globl = ops.Const(0) with ops.task_instance_init(): local = ops.Const(0) with ops.loop(100): ops.Copy(globl, local) with ops.task_instance_exit(): ops.Add([globl, local], [globl]) with ops.task_exit(): ops.Mul([globl, globl], [globl]) plan = core.Plan(task.get_step()) test_plan = core.Plan.create_from_proto(plan.Proto()) self.assertEqual(len(plan.Steps()), 1) self.assertEqual(len(test_plan.Steps()), 1) self.assertEqual(len(plan.Proto().network), 9) self.assertEqual(len(test_plan.Proto().network), 9) self.assertEqual(len(plan.Proto().execution_step), 1) self.assertEqual(len(test_plan.Proto().execution_step), 1) self.assertEqual(plan.Steps()[0].Name(), test_plan.Steps()[0].Name()) self.assertEqual(len(plan.Nets()), len(test_plan.Nets())) for idx in range(0, len(plan.Nets())): # When we create Net for test_plan, we will end up with new Net # name with postfix. net_1 = plan.Nets()[idx] net_2 = test_plan.Nets()[idx] trim_size = len(net_1.Name()) self.assertEqual(net_1.Name(), net_2.Name()[:trim_size])
def test_multi_instance(self): NUM_INSTANCES = 10 NUM_ITERS = 15 with TaskGroup() as tg: with Task(num_instances=NUM_INSTANCES): with ops.task_init(): counter1 = ops.CreateCounter([], ['global_counter']) counter2 = ops.CreateCounter([], ['global_counter2']) counter3 = ops.CreateCounter([], ['global_counter3']) # both task_counter and local_counter should be thread local with ops.task_instance_init(): task_counter = ops.CreateCounter([], ['task_counter']) local_counter = ops.CreateCounter([], ['local_counter']) with ops.loop(NUM_ITERS): ops.CountUp(counter1) ops.CountUp(task_counter) ops.CountUp(local_counter) # gather sum of squares of local counters to make sure that # each local counter counted exactly up to NUM_ITERS, and # that there was no false sharing of counter instances. with ops.task_instance_exit(): count2 = ops.RetrieveCount(task_counter) with ops.loop(ops.Mul([count2, count2])): ops.CountUp(counter2) # This should have the same effect as the above count3 = ops.RetrieveCount(local_counter) with ops.loop(ops.Mul([count3, count3])): ops.CountUp(counter3) # The code below will only run once with ops.task_exit(): total1 = final_output(ops.RetrieveCount(counter1)) total2 = final_output(ops.RetrieveCount(counter2)) total3 = final_output(ops.RetrieveCount(counter3)) with LocalSession() as session: session.run(tg) self.assertEquals(total1.fetch(), NUM_INSTANCES * NUM_ITERS) self.assertEquals(total2.fetch(), NUM_INSTANCES * (NUM_ITERS**2)) self.assertEquals(total3.fetch(), NUM_INSTANCES * (NUM_ITERS**2))
def example_loop(): with Task(): total = ops.Const(0) total_large = ops.Const(0) total_small = ops.Const(0) total_tiny = ops.Const(0) with ops.loop(10) as loop: outer = ops.Mul([loop.iter(), ops.Const(10)]) with ops.loop(loop.iter()) as inner: val = ops.Add([outer, inner.iter()]) with ops.If(ops.GE([val, ops.Const(80)])) as c: ops.Add([total_large, val], [total_large]) with c.Elif(ops.GE([val, ops.Const(50)])) as c: ops.Add([total_small, val], [total_small]) with c.Else(): ops.Add([total_tiny, val], [total_tiny]) ops.Add([total, val], total)
def _actual_loop(self): total = ops.Const(0) total_large = ops.Const(0) total_small = ops.Const(0) total_tiny = ops.Const(0) with ops.loop(10) as loop: outer = ops.Mul([loop.iter(), ops.Const(10)]) with ops.loop(loop.iter()) as inner: val = ops.Add([outer, inner.iter()]) with ops.If(ops.GE([val, ops.Const(80)])) as c: ops.Add([total_large, val], [total_large]) with c.Elif(ops.GE([val, ops.Const(50)])) as c: ops.Add([total_small, val], [total_small]) with c.Else(): ops.Add([total_tiny, val], [total_tiny]) ops.Add([total, val], total) return map(final_output, (total, total_large, total_small, total_tiny))