from inspect import currentframe, getframeinfo import unittest import numpy as np from caffe2.proto import caffe2_pb2 from caffe2.python import core, workspace, schema, test_util from caffe2.python.task import Node, Task class TestScopes(test_util.TestCase): def testBlobReferenceIsIndependentFromNameScope(self): blob_v = core.BlobReference("v") with core.NameScope("foo"): blob_w = core.BlobReference("w") with core.NameScope("bar"): blob_x = core.BlobReference("x") self.assertEqual(str(blob_v), "v") self.assertEqual(str(blob_w), "w") self.assertEqual(str(blob_x), "x") def testNameScopeWithOp(self): global_x = core.BlobReference("x") global_y = core.BlobReference("y") with core.NameScope("foo"): # Raw strings should have namescope prepended. op = core.CreateOperator("Relu", "x", "y") self.assertEqual(len(op.input), 1) self.assertEqual(op.input[0], "foo/x") self.assertEqual(len(op.output), 1) self.assertEqual(op.output[0], "foo/y") # BlobReferences should not. op = core.CreateOperator("Relu", global_x, global_y) self.assertEqual(len(op.input), 1) self.assertEqual(op.input[0], "x") self.assertEqual(len(op.output), 1) self.assertEqual(op.output[0], "y") def testNameScopeWithReset(self): with core.NameScope("foo"): # foo/ op = core.CreateOperator("Relu", "x", "y") self.assertEqual(len(op.input), 1) self.assertEqual(op.input[0], "foo/x") self.assertEqual(len(op.output), 1) self.assertEqual(op.output[0], "foo/y") with core.NameScope("bar"): # foo/bar/ op = core.CreateOperator("Relu", "x", "y") self.assertEqual(len(op.input), 1) self.assertEqual(op.input[0], "foo/bar/x") self.assertEqual(len(op.output), 1) self.assertEqual(op.output[0], "foo/bar/y") # Back to foo/ op = core.CreateOperator("Relu", "x", "y") self.assertEqual(len(op.input), 1) self.assertEqual(op.input[0], "foo/x") self.assertEqual(len(op.output), 1) self.assertEqual(op.output[0], "foo/y") with core.NameScope("bar", reset=True): # bar/ op = core.CreateOperator("Relu", "x", "y") self.assertEqual(len(op.input), 1) self.assertEqual(op.input[0], "bar/x") self.assertEqual(len(op.output), 1) self.assertEqual(op.output[0], "bar/y") # Back to foo/ op = core.CreateOperator("Relu", "x", "y") self.assertEqual(len(op.input), 1) self.assertEqual(op.input[0], "foo/x") self.assertEqual(len(op.output), 1) self.assertEqual(op.output[0], "foo/y") def testDeviceScope(self): # No device op = core.CreateOperator("Relu", "x", "y") self.assertFalse(op.HasField('device_option')) # explicitly setting a device device_option = caffe2_pb2.DeviceOption() device_option.device_type = workspace.GpuDeviceType device_option.device_id = 1 op = core.CreateOperator("Relu", "x", "y", device_option=device_option) self.assertTrue(op.HasField('device_option')) self.assertEqual(op.device_option.device_type, workspace.GpuDeviceType) self.assertEqual(op.device_option.device_id, 1) with core.DeviceScope(device_option): # from device scope op = core.CreateOperator("Relu", "x", "y") self.assertTrue(op.HasField('device_option')) self.assertEqual(op.device_option.device_type, workspace.GpuDeviceType) self.assertEqual(op.device_option.device_id, 1) # from an overridden device option override_device = caffe2_pb2.DeviceOption() override_device.device_type = caffe2_pb2.CPU op = core.CreateOperator( "Relu", "x", "y", device_option=override_device) self.assertTrue(op.HasField('device_option')) self.assertEqual(op.device_option.device_type, caffe2_pb2.CPU) # back from normal: no device op = core.CreateOperator("Relu", "x", "y") self.assertFalse(op.HasField('device_option')) device_option = caffe2_pb2.DeviceOption() def testNameAndDeviceScopeTogether(self): device_option = caffe2_pb2.DeviceOption() device_option.device_type = workspace.GpuDeviceType device_option.device_id = 1 with core.DeviceScope(device_option): with core.NameScope("foo"): op = core.CreateOperator("Relu", "x", "y") self.assertTrue(op.HasField('device_option')) self.assertEqual(op.device_option.device_type, workspace.GpuDeviceType) self.assertEqual(op.device_option.device_id, 1) self.assertEqual(len(op.input), 1) self.assertEqual(op.input[0], "foo/x") self.assertEqual(len(op.output), 1) self.assertEqual(op.output[0], "foo/y") class TestCloneNet(test_util.TestCase): def testPartialClone(self): params = core.Net('params') p1 = params.ConstantFill([], ['p1']) workspace.CreateNet(params) workspace.RunNetOnce(params) n = core.Net('original') a1 = n.AddExternalInput('a1') a2 = n.AddExternalInput('a2') b1, b2 = n.Concat([a1, a2], ['b1', 'b2'], axis=0) c1 = n.Sum([b1, p1], ['c1']) c2 = n.Sum([b2], ['c2']) d = n.Sum([c1, c2], ['d']) # test that gradient ops are ignored when partial-cloning n.AddGradientOperators([d]) # test some in-place ops k = n.Sum([p1], ['k']) e = n.Sum([d], ['e']) e = n.Sum([e, k], [e]) e = n.Sum([e], [e]) f = n.Sum(e, ['f']) def net_assert(net, num_ops, inputs, outputs, internals): self.assertEqual(len(net.Proto().op), num_ops) self.assertEqual(set(net.Proto().external_input), inputs) self.assertEqual(set(net.Proto().external_output), outputs) all_blobs = set(net.Proto().external_input) all_blobs |= set(net.Proto().external_output) for op in net.Proto().op: all_blobs |= set(op.input) | set(op.output) self.assertEqual(all_blobs, inputs | outputs | internals) # create net to make sure its valid for input in inputs: workspace.FeedBlob(input, np.array([])) workspace.CreateNet(net) n2, (d22, ) = n.ClonePartial('f1', {a1: 'a11', a2: 'a22'}, [d]) net_assert( n2, 4, {'p1', 'a11', 'a22'}, {'f1/d'}, {'f1/b1', 'f1/b2', 'f1/c1', 'f1/c2', 'p1'}) self.assertTrue(isinstance(d22, core.BlobReference)) self.assertEqual(d22.Net(), n2) self.assertEqual(str(d22), 'f1/d') n3, (d22, ) = n.ClonePartial('f2', [b1, b2], [d]) net_assert( n3, 3, {'p1', 'b1', 'b2'}, {'f2/d'}, {'f2/c1', 'f2/c2', 'p1'}) self.assertEqual(str(d22), 'f2/d') n4, (c22, ) = n.ClonePartial('f3', [b1], [c1]) net_assert(n4, 1, {'p1', 'b1'}, {'f3/c1'}, {'p1'}) self.assertEqual(str(c22), 'f3/c1') n5, (c11, c22) = n.ClonePartial('f4', [b1, b2], [c1, c2]) net_assert(n5, 2, {'p1', 'b1', 'b2'}, {'f4/c1', 'f4/c2'}, {'p1'}) self.assertEqual(str(c11), 'f4/c1') self.assertEqual(str(c22), 'f4/c2') with self.assertRaises(AssertionError): n.ClonePartial('f4', [a1, a2, c2], [d]) n6, (e22, ) = n.ClonePartial('f5', [d], [e]) net_assert(n6, 4, {'p1', 'd'}, {'f5/e'}, {'f5/k', 'p1'}) self.assertEqual(str(e22), 'f5/e') n8, (e22, f22) = n.ClonePartial('f7', [d], [e, f]) net_assert(n8, 5, {'p1', 'd'}, {'f7/e', 'f7/f'}, {'p1', 'f7/k'}) self.assertEqual(str(e22), 'f7/e') self.assertEqual(str(f22), 'f7/f') params._CheckLookupTables() n._CheckLookupTables() def test_mask_clone_update_external_list(self): n = core.Net('original') a1 = n.AddExternalInput('a1') a2 = n.AddExternalInput('a2') p1 = 'p1' b1, b2 = n.Concat([a1, a2], ['b1', 'b2'], axis=0) c1 = n.Sum([b1, p1], ['c1']) c2 = n.Sum([b2], ['c2']) n.Sum([c1, c2], ['d']) new_net = n.Clone( "new", op_id_mask=[0, 1], keep_schema=True, update_external_list=True) self.assertEqual( sorted(map(str, new_net.external_inputs)), ["a1", "a2", "p1"], "external input not matched", ) self.assertEqual( sorted(map(str, new_net.external_outputs)), ["b2", "c1"], "external output not matched", ) new_net = n.Clone( "new2", op_id_mask=[2, 3], keep_schema=True, update_external_list=True) self.assertEqual( sorted(map(str, new_net.external_inputs)), ["b2", "c1"], "external input not matched", ) self.assertEqual( sorted(map(str, new_net.external_outputs)), ["d"], "external output not matched", ) def test_control_op_remap(self): # Subnets under If/AsyncIf operators should get name remapping when cloned n = core.Net("original") then_net = core.Net("a") then_net.FC(["inputA"], "fc_a") else_net = core.Net("b") else_net.FC(["inputB"], "fc_b") n.If( inputs=[], outputs=[], then_net=then_net.Proto(), else_net=else_net.Proto(), ) copied = n.Clone("copied", blob_remap={"inputA": "inputX"}) if_op = copied._net.op[0] self.assertEqual(if_op.arg[0].n.op[0].input, ["inputX"]) self.assertEqual(if_op.arg[1].n.op[0].input, ["inputB"]) class TestExternalInputs(test_util.TestCase): def testAddExternalInputShouldRaiseIfDuplicate(self): net = core.Net("test") net.AddExternalInput( schema.Struct(("x", schema.Scalar(np.float))), ) with self.assertRaises(AssertionError): net.AddExternalInput( schema.Struct(("x", schema.Scalar(np.float))), ) def testAddExternalInputShouldRaiseIfDuplicateInSameCall(self): net = core.Net("test") with self.assertRaises(AssertionError): net.AddExternalInput( schema.Struct(("x", schema.Scalar(np.float))), schema.Struct(("x", schema.Scalar(np.float))), ) def testSetInputRecordWithBlobs(self): net = core.Net("test") record = schema.NewRecord(net, schema.Struct( ("x", schema.Scalar(np.float)), )) input_record = net.set_input_record(record) self.assertTrue(net.BlobIsDefined(input_record.x())) self.assertIn(input_record.x(), net.external_inputs) def testSetInputRecordWithoutBlobs(self): net = core.Net("test") record = schema.Struct(("x", schema.Scalar(np.float))) input_record = net.set_input_record(record) self.assertTrue(net.BlobIsDefined(input_record.x())) self.assertIn(input_record.x(), net.external_inputs) class TestCreateOperator(test_util.TestCase): def testCreate(self): device_option = caffe2_pb2.DeviceOption() device_option.device_type = workspace.GpuDeviceType device_option.device_id = 1 op = core.CreateOperator( "Ludicrous", "x", "y", name="ludicrous", control_input="z", device_option=device_option, engine="WARP", arg1=1, arg2="2", arg3=[1, 2, 3]) self.assertEqual(op.type, "Ludicrous") self.assertEqual(op.name, "ludicrous") self.assertEqual(op.engine, "WARP") self.assertEqual(len(op.input), 1) self.assertEqual(op.input[0], "x") self.assertEqual(len(op.output), 1) self.assertEqual(op.output[0], "y") self.assertEqual(len(op.control_input), 1) self.assertEqual(op.control_input[0], "z") self.assertTrue(op.HasField('device_option')) self.assertEqual(op.device_option.device_type, workspace.GpuDeviceType) self.assertEqual(op.device_option.device_id, 1) self.assertEqual(len(op.arg), 3) # can't guarantee ordering of kwargs, so generate a set of args # to test with arg_map = {} for arg in op.arg: arg_map[arg.name] = arg # Check all elements exist that should self.assertEqual("arg1" in arg_map, True) self.assertEqual("arg2" in arg_map, True) self.assertEqual("arg3" in arg_map, True) # Now test that all args were initialized correctly self.assertEqual(arg_map["arg1"].i, 1) self.assertEqual(arg_map["arg2"].s, b"2") self.assertEqual(list(arg_map["arg3"].ints), [1, 2, 3]) class TestAutoNaming(test_util.TestCase): def assertOperatorListEqual(self, operatorDefList1, operatorDefList2): for op in operatorDefList1: op.debug_info = "" for op in operatorDefList2: op.debug_info = "" self.assertEqual(operatorDefList1, operatorDefList2) """ Test that operators are named with different names, and that automatically named blob names don't clash intra or inter networks. """ def test_next_blob(self): def create_net(): net = core.Net('net') with core.NameScope('foo'): net.Add(['a', 'b'], net.NextScopedBlob('ab')) net.Add(['c', 'd'], net.NextBlob('cd')) return net net_a = create_net() net_b = create_net() # created net proto is predicatable. self.assertOperatorListEqual(net_a.Proto().op, net_b.Proto().op) self.assertEqual(net_a.Proto().op[0].output[0], 'foo/ab') self.assertEqual(net_a.Proto().op[1].output[0], 'cd') net_c = core.Net('net') # different calls return different blob names self.assertNotEqual(str(net_c.NextBlob('b')), str(net_c.NextBlob('b'))) def test_auto_naming(self): a = core.Net('net') b = core.Net('net') self.assertNotEqual(a.Proto().name, b.Proto().name) a_in1 = a.AddExternalInput('a') b_in1 = b.AddExternalInput('b') all_outputs_single = [] all_outputs_list = [] def add_ops(): all_outputs_single.append(a.Sum([a_in1, a_in1])) all_outputs_single.append(a.Sum([a_in1, a_in1])) all_outputs_single.append(b.Sum([b_in1, b_in1])) all_outputs_single.append(b.Sum([b_in1, b_in1])) all_outputs_list.append(a.Sum([a_in1, a_in1], outputs=2)) all_outputs_list.append(a.Sum([a_in1, a_in1], outputs=2)) all_outputs_list.append(b.Sum([b_in1, b_in1], outputs=2)) all_outputs_list.append(b.Sum([b_in1, b_in1], outputs=2)) add_ops() with core.NameScope('n1'): add_ops() # Force reset of lookup tables a.Proto().name with core.NameScope('n2'): add_ops() all_outputs = [] for s in all_outputs_single: all_outputs.append(str(s)) for l in all_outputs_list: for o in l: all_outputs.append(str(o)) for i, o1 in enumerate(all_outputs): for j, o2 in enumerate(all_outputs): if i != j: self.assertNotEqual(str(o1), str(o2)) a._CheckLookupTables() b._CheckLookupTables() class TestAppendNet(test_util.TestCase): def test_external_inputs_merged_correctly(self): netA = core.Net("A") netA.Sum(["in1", "in2"], ["sum1"]) self.assertTrue("in1" in netA.external_inputs) netB = core.Net("B") netB.Sum(["in3", "in4"], ["in1"]) netB.AppendNet(netA) self.assertFalse("in1" in netB.external_inputs) def test_external_inputs_merged_correctlyB(self): netA = core.Net("A") netA.Sum(["in1", "in2"], ["sum1"]) self.assertTrue("in1" in netA.external_inputs) netB = core.Net("B") netB.Sum(["in3", "in4"], ["in1"]) netA.AppendNet(netB) # note different order than in prev test self.assertTrue("in1" in netA.external_inputs) class TestExtractPredictorNet(test_util.TestCase): @unittest.skipIf('ImageInput' not in workspace.RegisteredOperators(), "Needs OpenCV") def test_extract_simple(self): from caffe2.python import brew from caffe2.python.model_helper import ModelHelper, ExtractPredictorNet model = ModelHelper(name="test", arg_scope={'order': 'NCHW'}) [data, label] = brew.image_input( model, "reader", ["xx/data", "label"], is_test=1, ) cnv = brew.conv(model, data, 'cnv', 32, 32, 4) a = brew.fc(model, cnv, 'a', 100, 200) pred = brew.fc(model, a, 'pred', 200, 5) brew.softmax(model, [pred, label], "softmax") (predict_net, export_blobs) = ExtractPredictorNet( net_proto=model.net.Proto(), input_blobs=["xx/data"], output_blobs=["pred"], renames={"xx/data": "image"}, ) export_blobs = set(export_blobs) ops = list(predict_net.Proto().op) for op in ops: self.assertFalse(op.type == "Softmax") self.assertFalse("xx/data" in op.input) # Note: image input should not be included self.assertEquals(ops[0].type, "Conv") self.assertEquals(ops[1].type, "FC") self.assertEquals(ops[2].type, "FC") self.assertEquals(len(ops), 3) # test rename happened self.assertEquals(ops[0].input[0], "image") # Check export blobs self.assertTrue("image" not in export_blobs) self.assertTrue("xx/data" not in export_blobs) self.assertEqual(set([str(p) for p in model.params]), export_blobs) # Check external inputs/outputs self.assertTrue("image" in predict_net.Proto().external_input) self.assertEquals(set(["pred"]), set(predict_net.Proto().external_output)) self.assertEqual( set(predict_net.Proto().external_input) - set([str(p) for p in model.params]), set(["image"]) ) class TestOperatorTraceback(test_util.TestCase): def op_name_check(self, net, cf, line, func): net.PopulateProtoWithFileName() filename = getframeinfo(cf).filename self.assertEqual(net.Proto().op[0].name, '{}:{}:{}'.format( filename, line, func)) def test_operator_constructor_traceback(self): net = core.Net("test") a, b = net.AddExternalInput("a", "b") net.Mul([a, b], "c"); cf = currentframe(); line = cf.f_lineno func = cf.f_code.co_name with self.assertRaises(Exception): workspace.RunNetOnce(net) with self.assertRaises(Exception): workspace.CreateNet(net) self.op_name_check(net, cf, line, func) def test_operator_runtime_traceback(self): net = core.Net("test") a = net.AddExternalInput("a") workspace.blobs[a] = np.array([1, 2, 3], dtype=np.float32) net.Split(a, ["b", "c"], axis=0); cf = currentframe(); line = cf.f_lineno func = cf.f_code.co_name with self.assertRaises(Exception): workspace.RunNetOnce(net) workspace.CreateNet(net) with self.assertRaises(Exception): workspace.RunNet(net) self.op_name_check(net, cf, line, func) def test_c_workspace_constructor(self): net = core.Net("test") a, b = net.AddExternalInput("a", "b") net.Mul([a, b], "c"); cf = currentframe(); line = cf.f_lineno func = cf.f_code.co_name ws = workspace.C.Workspace() with self.assertRaises(Exception): ws.run(net) with self.assertRaises(Exception): ws.create_net(net) self.op_name_check(net, cf, line, func) def test_c_workspace_runtime(self): net = core.Net("test") a = net.AddExternalInput("a") net.Split(a, ["b", "c"], axis=0); cf = currentframe(); line = cf.f_lineno func = cf.f_code.co_name ws = workspace.C.Workspace() ws.create_blob(str(a)).feed(np.array([1, 2, 3], dtype=np.float32)) ws.create_net(net) with self.assertRaises(Exception): ws.run(net) self.op_name_check(net, cf, line, func) def test_async_exception_handling(self): net = core.Net("test") net.Proto().type = 'dag' # this runs operators on background threads a = net.AddExternalInput("a") net.Split(a, ["b", "c"], axis=0); cf = currentframe(); line = cf.f_lineno func = cf.f_code.co_name workspace.FeedBlob(a, np.array([1, 2, 3], dtype=np.float32)) with self.assertRaises(Exception) as enforceNotMet: workspace.RunNetOnce(net) self.assertIn('enforce fail', str(enforceNotMet.exception)) self.op_name_check(net, cf, line, func) class TestCreatePlan(test_util.TestCase): 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]) class TestOpRegistryKey(test_util.TestCase): def test_is_operator(self): self.assertTrue(core.IsOperator('Relu')) self.assertFalse(core.IsOperator('NOEXIST')) def test_is_operator_with_engine(self): self.assertTrue(core.IsOperatorWithEngine('Relu', 'DEFAULT')) self.assertFalse(core.IsOperatorWithEngine('Relu', 'NOEXIST')) class TestDeviceOption(test_util.TestCase): def test_check_equal_node_name(self): opt1 = core.DeviceOption(0) opt2 = core.DeviceOption(0) self.assertTrue(core.device_option_equal(opt1, opt2)) opt2.node_name = 'test' self.assertTrue(core.device_option_equal(opt1, opt2)) self.assertFalse(core.device_option_equal(opt1, opt2, ignore_node_name=False)) opt1.node_name = 'test' self.assertTrue(core.device_option_equal(opt1, opt2, ignore_node_name=False)) def test_check_equal_default_value(self): opt1 = caffe2_pb2.DeviceOption() opt2 = caffe2_pb2.DeviceOption() opt1.device_type = 0 self.assertTrue(core.device_option_equal(opt1, opt2)) opt1.device_id = 5 # opt1 still is on CPU, so the options should be equal self.assertTrue(core.device_option_equal(opt1, opt2)) opt2.device_type = 0 self.assertTrue(core.device_option_equal(opt1, opt2)) opt1.device_type = 1 self.assertFalse(core.device_option_equal(opt1, opt2)) class TestInferDeviceCpuOnly(test_util.TestCase): def test_inject_copy(self): ''' Test inject cross device copies - this is a no-op on CPU only devices. ''' send_node = 'node:0' recv_node = 'node:1' # Using placeholder ops for send/recv. Placeholder ops are # decorator/fake ops that don't have operator schema. placeholder_send = 'Placeholder:Dummy:Send' placeholder_recv = 'Placeholder:Dummy:Recv' # init_net. init_net = core.Net("init_net") with core.DeviceScope(0, node_name=send_node): init_net.XavierFill([], 'fc_w', shape=[10, 100]) init_net.ConstantFill([], 'fc_b', shape=[10, ]) # train_net. train_net = core.Net("train_net") train_net.Proto().external_input.extend(['fc_w', 'fc_b']) with core.DeviceScope(0, node_name=send_node): op = core.CreateOperator( placeholder_send, ["fc_w", 'fc_b'], [], dst_node=recv_node) train_net.Proto().op.extend([op]) with core.DeviceScope(0, node_name=recv_node): # Let's rename the recv blob i.e. fc_w -> fc_w_recv. op = core.CreateOperator( placeholder_recv, [], ['fc_w_recv', 'fc_b'], src_node=send_node) train_net.Proto().op.extend([op]) train_net.FC(["data", 'fc_w_recv', 'fc_b'], "fc1") # Inject cross device copies. init_net, x_dev_state = core.InjectCrossDeviceCopies( init_net, placeHolderOps=[placeholder_send, placeholder_recv]) train_net, x_dev_state = core.InjectCrossDeviceCopies( train_net, x_dev_state, placeHolderOps=[placeholder_send, placeholder_recv]) # Verify: No Copy operators should be injected since it is CPU only. op = train_net.Proto().op[0] self.assertEqual(op.type, placeholder_send) self.assertEqual(op.device_option.device_type, 0) self.assertEqual(op.input[0], "fc_w") self.assertEqual(op.input[1], "fc_b") op = train_net.Proto().op[1] self.assertEqual(op.type, placeholder_recv) self.assertEqual(op.device_option.device_type, 0) self.assertEqual(op.output[0], "fc_w_recv") self.assertEqual(op.output[1], "fc_b") op = train_net.Proto().op[2] self.assertEqual(op.type, "FC") self.assertEqual(op.device_option.device_type, 0) self.assertEqual(op.input[1], "fc_w_recv") self.assertEqual(op.input[2], "fc_b") @unittest.skipIf(not workspace.has_gpu_support, 'No GPU support') class TestInferDevice(test_util.TestCase): def setUp(self): device_option = caffe2_pb2.DeviceOption() device_option.device_type = workspace.GpuDeviceType device_option.device_id = 1 self.gpu_option = device_option self.cpu_option = caffe2_pb2.DeviceOption() def _test_op( self, op_name, in_option, out_option, op_option=None, inputs=None, outputs=None ): op_option = self.gpu_option if not op_option else op_option inputs = ["blob_1"] if not inputs else inputs outputs = ["blob_2"] if not outputs else outputs with core.DeviceScope(op_option): op = core.CreateOperator(op_name, inputs, outputs) input_dev, output_dev = core.InferOpBlobDevices(op) if isinstance(in_option, list): assert len(in_option) == len(input_dev), \ 'Length of input device option should match' \ '{} vs. {}'.format(in_option, input_dev) for in_dev, in_opt in zip(input_dev, in_option): self.assertEqual(in_dev, in_opt) else: for in_dev in input_dev: self.assertEqual(in_dev, in_option) if isinstance(out_option, list): assert len(out_option) == len(output_dev), \ 'Length of output device option should match' \ '{} vs. {}'.format(out_option, output_dev) for out_dev, out_opt in zip(output_dev, out_option): self.assertEqual(out_dev, out_opt) else: for out_dev in output_dev: self.assertEqual(out_dev, out_option) def test_infer_device(self): self._test_op( "FC", self.gpu_option, self.gpu_option, op_option=self.gpu_option, inputs=["data", "fc_w", "fc_b"], outputs=["fc_1"] ) def test_infer_device_split_by_lengths(self): self._test_op( "SplitByLengths", [self.gpu_option, self.cpu_option], self.gpu_option, op_option=self.gpu_option, inputs=["data", "fc_w"], outputs=["fc_1"] ) def test_infer_device_adam(self): in_options = [self.gpu_option] * 6 in_options[5] = self.cpu_option out_options = [self.gpu_option] * 4 self._test_op( "Adam", in_options, out_options, op_option=self.gpu_option, inputs=["param", "moment_1", "moment_2", "grad", "lr", "iter"], outputs=["output_param", "output_moment_1", "output_moment_2", "output_grad"] ) def test_infer_device_cross_device(self): self._test_op("CopyGPUToCPU", self.gpu_option, self.cpu_option) self._test_op("CopyCPUToGPU", self.cpu_option, self.gpu_option) self._test_op("CopyFromCPUInput", self.cpu_option, self.gpu_option) self._test_op( "CopyFromCPUInput", self.cpu_option, self.cpu_option, op_option=self.cpu_option ) def test_device_inference_function(self): # ConcatOp. op_option = self.gpu_option with core.DeviceScope(op_option): op = core.CreateOperator( 'Concat', ['X_{}'.format(i) for i in range(4)], ['concat_result', 'split_info'], axis=1) input_dev, output_dev = core.InferOpBlobDevices(op) # 2nd output's type is CPU irrespective of Concat op's device option. self.assertEqual(output_dev[1], self.cpu_option) #SplitOp. op_option = self.gpu_option with core.DeviceScope(op_option): op = core.CreateOperator( 'Split', ['input', 'split'], ['X_{}'.format(i) for i in range(4)], axis=0) input_dev, output_dev = core.InferOpBlobDevices(op) # 2nd input's type is CPU irrespective of Split op's device option. self.assertEqual(input_dev[1], self.cpu_option) def test_inject_copy(self): net = core.Net("test") init_net = core.Net("init") device_option = caffe2_pb2.DeviceOption() device_option.device_type = workspace.GpuDeviceType device_option.device_id = 1 weight = init_net.XavierFill([], 'fc_w', shape=[10, 100]) bias = init_net.ConstantFill([], 'fc_b', shape=[10, ]) with core.DeviceScope(device_option): net.FC(["data", weight, bias], "fc1") _, blob_to_device = core.InjectCrossDeviceCopies(init_net) new_net, blob_to_device = core.InjectCrossDeviceCopies( net, blob_to_device ) op = new_net._net.op[-1] self.assertEqual(op.type, "FC") self.assertEqual(op.input[0], "data_gpu_1") self.assertEqual(op.input[1], "fc_w_gpu_1") self.assertEqual(op.input[2], "fc_b_gpu_1") self.assertEqual(op.device_option.device_type, workspace.GpuDeviceType) self.assertEqual(op.device_option.device_id, 1) self.assertEqual(new_net._net.op[-2].type, "CopyCPUToGPU") self.assertEqual(new_net._net.op[0].type, "CopyCPUToGPU") self.assertNotEqual(blob_to_device["fc_w"], device_option) def test_cross_nets(self): net = core.Net("test") init_net = core.Net("init") device_option = caffe2_pb2.DeviceOption() device_option.device_type = workspace.GpuDeviceType device_option.device_id = 1 weight = init_net.XavierFill([], 'fc_w', shape=[10, 100]) bias = init_net.ConstantFill([], 'fc_b', shape=[10, ]) const = init_net.ConstantFill([], 'const', shape=[], value=1.) with core.DeviceScope(device_option): const = init_net.Add([const, const], [const]) fc_out = net.FC(["data", weight, bias], "fc1") net.Add([fc_out, const], [fc_out]) data_remap = {'data': device_option} nets, _ = core.InjectDeviceCopiesAmongNets( [init_net, net], blob_to_device_init=data_remap ) op = nets[1]._net.op[0] self.assertEqual(op.type, "CopyCPUToGPU") self.assertEqual(op.device_option.device_type, workspace.GpuDeviceType) self.assertEqual(op.device_option.device_id, 1) self.assertEqual(op.output[0], "fc_w_gpu_1") op = nets[1]._net.op[1] self.assertEqual(op.type, "CopyCPUToGPU") self.assertEqual(op.device_option.device_type, workspace.GpuDeviceType) self.assertEqual(op.device_option.device_id, 1) self.assertEqual(op.output[0], "fc_b_gpu_1") op = nets[1]._net.op[2] self.assertEqual(op.type, "FC") self.assertEqual(op.input[0], "data") self.assertEqual(op.input[1], "fc_w_gpu_1") self.assertEqual(op.input[2], "fc_b_gpu_1") self.assertEqual(op.device_option.device_type, workspace.GpuDeviceType) self.assertEqual(op.device_option.device_id, 1) op = nets[1]._net.op[3] self.assertEqual(op.type, "Add") self.assertEqual(op.input[0], "fc1") self.assertEqual(op.input[1], "const_gpu_1") # check that moved blob is in input to the new net for c in ["data", "fc_w", "fc_b", "const_gpu_1"]: self.assertTrue(c in nets[1]._net.external_input) """ For reference, net.Proto() should be like: name: "" op { input: "fc_w" output: "fc_w_gpu_1" name: "" type: "CopyCPUToGPU" device_option { device_type: 1 device_id: 1 } } op { input: "fc_b" output: "fc_b_gpu_1" name: "" type: "CopyCPUToGPU" device_option { device_type: 1 device_id: 1 } } op { input: "data" input: "fc_w_gpu_1" input: "fc_b_gpu_1" output: "fc1" name: "" type: "FC" device_option { device_type: 1 device_id: 1 } } op { input: "fc1" input: "const_gpu_1" output: "fc1" name: "" type: "Add" device_option { device_type: 1 device_id: 1 } } external_input: "data" external_input: "fc_w" external_input: "fc_b" external_input: "const" external_input: "const_gpu_1" """ def test_cross_nets_no_change(self): net = core.Net("test") init_net = core.Net("init") device_option = caffe2_pb2.DeviceOption() device_option.device_type = workspace.GpuDeviceType device_option.device_id = 1 with core.DeviceScope(device_option): weight = init_net.XavierFill([], 'fc_w', shape=[10, 100]) bias = init_net.ConstantFill([], 'fc_b', shape=[10, ]) net.FC(["data", weight, bias], "fc1") data_remap = {'data': device_option} nets = core.InjectDeviceCopiesAmongNetsWithoutB2D( [init_net, net], blob_to_device_init=data_remap ) op = nets[1]._net.op[0] self.assertEqual(op.type, "FC") self.assertEqual(op.input[0], "data") self.assertEqual(op.input[1], "fc_w") self.assertEqual(op.input[2], "fc_b") self.assertEqual(op.device_option.device_type, workspace.GpuDeviceType) self.assertEqual(op.device_option.device_id, 1) """ For reference, net.Proto() should be like: name: "" op { input: "data" input: "fc_w" input: "fc_b" output: "fc1" name: "" type: "FC" device_option { device_type: 1 device_id: 1 } } external_input: "data" external_input: "fc_w" external_input: "fc_b" """ def test_inject_copy_multi_use(self): net = core.Net("test") device_option = caffe2_pb2.DeviceOption() device_option.device_type = workspace.GpuDeviceType device_option.device_id = 1 with core.DeviceScope(device_option): net.Relu("data", "relu1") net.Relu("data", "relu2") with core.DeviceScope(device_option): net.Relu("data", "relu3") net.Relu("data", "relu4") device_option.device_id = 0 with core.DeviceScope(device_option): net.Relu("data", "relu5") device_option.device_id = 1 with core.DeviceScope(device_option): net.Relu("data", "relu6") new_net, _ = core.InjectCrossDeviceCopies(net) op = new_net._net.op[0] self.assertEqual(op.type, "CopyCPUToGPU") self.assertEqual(op.device_option.device_type, workspace.GpuDeviceType) self.assertEqual(op.device_option.device_id, 1) self.assertEqual(op.output[0], "data_gpu_1") op = new_net._net.op[1] self.assertEqual(op.type, "Relu") self.assertEqual(op.device_option.device_type, workspace.GpuDeviceType) self.assertEqual(op.device_option.device_id, 1) self.assertEqual(op.output[0], "relu1") op = new_net._net.op[2] self.assertEqual(op.type, "Relu") self.assertEqual(op.device_option.device_type, 0) self.assertEqual(op.output[0], "relu2") op = new_net._net.op[3] self.assertEqual(op.type, "Relu") self.assertEqual(op.device_option.device_type, workspace.GpuDeviceType) self.assertEqual(op.device_option.device_id, 1) self.assertEqual(op.input[0], "data_gpu_1") self.assertEqual(op.output[0], "relu3") op = new_net._net.op[4] self.assertEqual(op.type, "Relu") self.assertEqual(op.device_option.device_type, 0) self.assertEqual(op.output[0], "relu4") op = new_net._net.op[5] self.assertEqual(op.type, "CopyCPUToGPU") self.assertEqual(op.device_option.device_type, workspace.GpuDeviceType) self.assertEqual(op.device_option.device_id, 0) self.assertEqual(op.output[0], "data_gpu_0") op = new_net._net.op[6] self.assertEqual(op.type, "Relu") self.assertEqual(op.device_option.device_type, workspace.GpuDeviceType) self.assertEqual(op.device_option.device_id, 0) self.assertEqual(op.input[0], "data_gpu_0") self.assertEqual(op.output[0], "relu5") op = new_net._net.op[7] self.assertEqual(op.type, "Relu") self.assertEqual(op.device_option.device_type, workspace.GpuDeviceType) self.assertEqual(op.device_option.device_id, 1) self.assertEqual(op.input[0], "data_gpu_1") self.assertEqual(op.output[0], "relu6") """ For reference, net.Proto() should be like: name: "" op { input: "data" output: "data_gpu_1" name: "" type: "CopyCPUToGPU" device_option { device_type: 1 device_id: 1 } } op { input: "data_gpu_1" output: "relu1" name: "" type: "Relu" device_option { device_type: 1 device_id: 1 } } op { input: "data" output: "relu2" name: "" type: "Relu" } op { input: "data_gpu_1" output: "relu3" name: "" type: "Relu" device_option { device_type: 1 device_id: 1 } } op { input: "data" output: "relu4" name: "" type: "Relu" } op { input: "data" output: "data_gpu_0" name: "" type: "CopyCPUToGPU" device_option { device_type: 1 device_id: 0 } } op { input: "data_gpu_0" output: "relu5" name: "" type: "Relu" device_option { device_type: 1 device_id: 0 } } op { input: "data_gpu_1" output: "relu6" name: "" type: "Relu" device_option { device_type: 1 device_id: 1 } } external_input: "data" """ def test_inject_copy_placeholder_ops(self): ''' Test inject cross device copies with placeholder ops. Placeholder ops are decorator/fake ops that don't have operator schema. ''' # Create CPU and GPU devices on 2 nodes. cpu_device = [] gpu_device = [] for i in range(0, 2): cpu_device.append(caffe2_pb2.DeviceOption()) cpu_device[i].node_name = 'node:' + str(i) gpu_device.append(caffe2_pb2.DeviceOption()) gpu_device[i].device_type = workspace.GpuDeviceType gpu_device[i].device_id = 0 gpu_device[i].node_name = 'node:' + str(i) send_node = 'node:0' recv_node = 'node:1' placeholder_send = 'Placeholder:Dummy:Send' placeholder_recv = 'Placeholder:Dummy:Recv' # init_net. init_net = core.Net("init_net") with core.DeviceScope(gpu_device[0]): weight = init_net.XavierFill([], 'fc_w', shape=[10, 100]) bias = init_net.ConstantFill([], 'fc_b', shape=[10, ]) with core.DeviceScope(cpu_device[0]): op = core.CreateOperator( placeholder_send, [weight, bias], [], dst_node=recv_node) init_net._net.op.extend([op]) # train_net train_net = core.Net("train_net") with core.DeviceScope(cpu_device[1]): # XXX. replace hardcoded op name. Move test to net_transforms. op = core.CreateOperator( placeholder_recv, [], [weight, bias], src_node=send_node) train_net._net.op.extend([op]) train_net.FC(["data", weight, bias], "fc1") # Inject cross device copies. init_net, x_dev_state = core.InjectCrossDeviceCopies( init_net, placeHolderOps=[placeholder_send, placeholder_recv]) train_net, x_dev_state = core.InjectCrossDeviceCopies( train_net, x_dev_state, placeHolderOps=[placeholder_send, placeholder_recv]) # Verify (init_net) op = init_net._net.op[2] self.assertEqual(op.type, "CopyGPUToCPU") self.assertEqual(op.device_option.device_type, workspace.GpuDeviceType) self.assertEqual(op.device_option.device_id, 0) self.assertEqual(op.output[0], "fc_w_cpu") op = init_net._net.op[3] self.assertEqual(op.type, "CopyGPUToCPU") self.assertEqual(op.device_option.device_type, workspace.GpuDeviceType) self.assertEqual(op.device_option.device_id, 0) self.assertEqual(op.output[0], "fc_b_cpu") op = init_net._net.op[4] self.assertEqual(op.type, placeholder_send) self.assertEqual(op.device_option.device_type, 0) self.assertEqual(op.input[0], "fc_w_cpu") self.assertEqual(op.input[1], "fc_b_cpu") # Verify (train_net) op = train_net._net.op[0] self.assertEqual(op.type, placeholder_recv) self.assertEqual(op.device_option.device_type, 0) self.assertEqual(op.output[0], "fc_w_cpu") self.assertEqual(op.output[1], "fc_b_cpu") op = train_net._net.op[3] self.assertEqual(op.type, "FC") self.assertEqual(op.device_option.device_type, 0) self.assertEqual(op.input[1], "fc_w_cpu") self.assertEqual(op.input[2], "fc_b_cpu") def test_blob_inplace(self): net = core.Net("test") device_option = caffe2_pb2.DeviceOption() device_option.device_type = workspace.GpuDeviceType device_option.device_id = 1 net.Adagrad(['param', 'moment', 'grad', 'lr'], ['param', 'moment']) with core.DeviceScope(device_option): net.Relu("param", "param_relu_no_sense") net, _ = core.InjectCrossDeviceCopies(net) op = net._net.op[1] self.assertEqual(op.type, 'CopyCPUToGPU') self.assertEqual(op.input[0], 'param') self.assertEqual(op.output[0], 'param_gpu_1') op = net._net.op[2] self.assertEqual(op.input[0], 'param_gpu_1') net.Relu('nonsense_input', 'moment') # should not raise inplace error core.InjectCrossDeviceCopies(net) with core.DeviceScope(device_option): net.Relu('nonsense_input_gpu', 'moment') with self.assertRaises(RuntimeError): core.InjectCrossDeviceCopies(net) class TestRerouteTensor(test_util.TestCase): def test_reroute_tensor(self): net = core.Net("reroute_tensor") net.Conv(["input", "w", "b"], "conv1") net.Relu(["conv1"], "conv1_relu") new_op = core.CreateOperator("SpatialBN", ["conv1", "scale", "bias", "mean", "var"], ["conv1_bn", "mean", "var", "saved_mean", "saved_var"]) # insert bn between conv and relu net.reroute_tensor("conv1", new_op, [net.Proto().op[1]]) self.assertEqual(new_op, net.Proto().op[1], "insertion failed") self.assertEqual(net.Proto().op[2].input[0], "conv1_bn", "reroute failed") class TestRunAllOnGPU(test_util.TestCase): def test_rnn_run_on_gpu(self): step_net = core.Net("step_net") step_net.Conv(["input_1", "w", "b"], "conv1") step_net.Relu(["conv1"], "input_1") net = core.Net("to_run_on_gpu") net.RecurrentNetwork(["input_1"], ["input_1"], step_net=step_net.Proto()) net.Relu(["input_1"], "input_relu") # check network structure before conversion net_proto = net.Proto() self.assertFalse(net_proto.HasField('device_option')) self.assertTrue(net_proto.op[0].arg[0].name == 'step_net') self.assertTrue(net_proto.op[0].arg[0].HasField('n')) self.assertFalse(net_proto.op[0].arg[0].n.HasField('device_option')) net.RunAllOnGPU(gpu_id=3, use_cudnn=True) # check that root net and rnn net got device_option attribute assigned self.assertTrue(net_proto.HasField('device_option')) self.assertEqual(net_proto.device_option.device_type, workspace.GpuDeviceType) self.assertEqual(net_proto.device_option.device_id, 3) self.assertTrue(net_proto.op[0].arg[0].n.HasField('device_option')) class TestConstructionFromProto(test_util.TestCase): def test_inplace_construction(self): # just create some random net n = core.Net('original') a1 = n.AddExternalInput('a1') a2 = n.AddExternalInput('a2') b1, b2 = n.Concat([a1, a2], ['b1', 'b2'], axis=0) c1 = n.Sum([b1, b1], ['c1']) c2 = n.Sum([b2], ['c2']) d = n.Sum([c1, c2], ['d']) proto = n.Proto() n_copied = core.Net(proto) n_moved = core.Net(proto, inplace=True) self.assertTrue(n_moved.Proto() is proto) self.assertTrue(n_copied.Proto() is not proto) proto.external_input.extend(['foo']) self.assertEqual(len(n_moved.Proto().external_input), len(proto.external_input)) self.assertEqual(len(n_copied.Proto().external_input), len(proto.external_input) - 1) if __name__ == '__main__': unittest.main()