def test_that_adding_gradient_scope_does_no_fancy_renaming(self):
# because it cannot create collisions
op = caffe2_pb2.OperatorDef()
op.name = 'foo_grad'
op.input.extend(['foo_grad', 'foo_grad_1'])
shapes = {'foo_grad': [1]}
blob_name_tracker = tb._get_blob_names([op])
tb._add_gradient_scope(shapes, blob_name_tracker, [op])
self.assertEqual(op.input[0], 'GRADIENTS/foo_grad')
self.assertEqual(op.input[1], 'GRADIENTS/foo_grad_1')
self.assertEqual(op.name, 'GRADIENTS/foo_grad')
self.assertEqual(len(shapes), 1)
self.assertEqual(shapes['GRADIENTS/foo_grad'], [1])
self.assertEqual(len(blob_name_tracker), 2)
self.assertEqual(blob_name_tracker['GRADIENTS/foo_grad'], 'foo_grad')
self.assertEqual(blob_name_tracker['GRADIENTS/foo_grad_1'],
'foo_grad_1')
def test_that_replacing_colons_gives_non_colliding_names(self):
# .. and update shapes
op = caffe2_pb2.OperatorDef()
op.name = 'foo:0'
op.input.extend(['foo:0', 'foo$0'])
shapes = {'foo:0': [1]}
blob_name_tracker = tb._get_blob_names([op])
tb._replace_colons(shapes, blob_name_tracker, [op], '$')
self.assertEqual(op.input[0], 'foo$0')
self.assertEqual(op.input[1], 'foo$0_1')
# Collision but blobs and op names are handled later by
# _fill_missing_operator_names.
self.assertEqual(op.name, 'foo$0')
self.assertEqual(len(shapes), 1)
self.assertEqual(shapes['foo$0'], [1])
self.assertEqual(len(blob_name_tracker), 2)
self.assertEqual(blob_name_tracker['foo$0'], 'foo:0')
self.assertEqual(blob_name_tracker['foo$0_1'], 'foo$0')
def test_that_auto_ssa_gives_non_colliding_names(self):
op1 = caffe2_pb2.OperatorDef()
op1.output.extend(['foo'])
op2 = caffe2_pb2.OperatorDef()
op2.input.extend(['foo'])
op2.output.extend(['foo'])
op2.output.extend(['foo_1'])
shapes = {'foo': [1], 'foo_1': [2]}
blob_name_tracker = tb._get_blob_names([op1, op2])
tb._convert_to_ssa(shapes, blob_name_tracker, [op1, op2])
self.assertEqual(op1.output[0], 'foo')
self.assertEqual(op2.input[0], 'foo')
self.assertEqual(op2.output[0], 'foo_1')
# Unfortunate name but we do not parse original `_` for now.
self.assertEqual(op2.output[1], 'foo_1_1')
self.assertEqual(len(shapes), 3)
self.assertEqual(shapes['foo'], [1])
self.assertEqual(shapes['foo_1'], [1])
self.assertEqual(shapes['foo_1_1'], [2])
self.assertEqual(len(blob_name_tracker), 3)
self.assertEqual(blob_name_tracker['foo'], 'foo')
self.assertEqual(blob_name_tracker['foo_1'], 'foo')
self.assertEqual(blob_name_tracker['foo_1_1'], 'foo_1')
def test_renaming_tensorflow_style(self):
# Construct some dummy operators here
# NOTE: '_w', '_bn', etc without the postfix '_' are only renamed when
# they are at the very end of the name.
# Test that '_w', '_w_' are renamed to '/weight', '/weight_', resp.
op1 = caffe2_pb2.OperatorDef()
op1.input.extend(['foo_w'])
op1.output.extend(['foo_w_2'])
# Test that '_bn', '_bn_' are renamed to '/batchnorm', '/batchnorm_',
# respectively.
op2 = caffe2_pb2.OperatorDef()
op2.input.extend(['foo_bn'])
op2.output.extend(['foo_bn_2'])
# Test that '_b', '_b_', are renamed to '/bias', '/bias_', resp.
op3 = caffe2_pb2.OperatorDef()
op3.input.extend(['foo_b'])
op3.output.extend(['foo_b_2'])
# Test that '_s', '_s_', are renamed to '/scale', '/scale_', resp.
op4 = caffe2_pb2.OperatorDef()
op4.input.extend(['foo_s'])
op4.output.extend(['foo_s_2'])
# Test that '_sum', '_sum_', are renamed to '/sum', '/sum_', resp.
op5 = caffe2_pb2.OperatorDef()
op5.input.extend(['foo_sum'])
op5.output.extend(['foo_sum_2'])
# Test that '_branch', '_branch_', are renamed to '/branch', '/branch_',
# respectively. Multiple inputs/outputs are also tested in this case.
op6 = caffe2_pb2.OperatorDef()
op6.input.extend(['foo_branch'])
op6.input.extend(['test_branch_2'])
op6.output.extend(['foo_branch_3'])
op6.output.extend(['test_branch4'])
shapes = {
'foo_w': [1], 'foo_w_2': [2], 'foo_bn': [3], 'foo_bn_2': [4],
'foo_b': [5], 'foo_b_2': [6], 'foo_s': [7], 'foo_s_2': [8],
'foo_sum': [9], 'foo_sum_2': [10], 'foo_branch': [11],
'test_branch_2': [12], 'foo_branch_3': [13], 'test_branch4': [14],
}
ops = [op1, op2, op3, op4, op5, op6]
blob_name_tracker = tb._get_blob_names(ops)
tb._rename_tensorflow_style(shapes, blob_name_tracker, ops)
# Testing that keys in blob name tracker were renamed correctly
self.assertEqual(blob_name_tracker['foo/weight'], 'foo_w')
self.assertEqual(blob_name_tracker['foo/weight_2'], 'foo_w_2')
self.assertEqual(blob_name_tracker['foo/batchnorm'], 'foo_bn')
self.assertEqual(blob_name_tracker['foo/batchnorm_2'], 'foo_bn_2')
self.assertEqual(blob_name_tracker['foo/bias'], 'foo_b')
self.assertEqual(blob_name_tracker['foo/bias_2'], 'foo_b_2')
self.assertEqual(blob_name_tracker['foo/scale'], 'foo_s')
self.assertEqual(blob_name_tracker['foo/scale_2'], 'foo_s_2')
self.assertEqual(blob_name_tracker['foo/sum'], 'foo_sum')
self.assertEqual(blob_name_tracker['foo/sum_2'], 'foo_sum_2')
self.assertEqual(blob_name_tracker['foo/branch'], 'foo_branch')
self.assertEqual(blob_name_tracker['test/branch_2'], 'test_branch_2')
self.assertEqual(blob_name_tracker['foo/branch_3'], 'foo_branch_3')
self.assertEqual(blob_name_tracker['test/branch4'], 'test_branch4')
# Testing that keys in shapes were renamed correctly
self.assertEqual(shapes['foo/weight'], [1])
self.assertEqual(shapes['foo/batchnorm_2'], [4])
self.assertEqual(shapes['foo/sum'], [9])
self.assertEqual(shapes['test/branch_2'], [12])
# Testing that the ops were renamed correctly
self.assertEqual(op1.input[0], 'foo/weight')
self.assertEqual(op1.output[0], 'foo/weight_2')
self.assertEqual(op2.input[0], 'foo/batchnorm')
self.assertEqual(op2.output[0], 'foo/batchnorm_2')
self.assertEqual(op3.input[0], 'foo/bias')
self.assertEqual(op3.output[0], 'foo/bias_2')
self.assertEqual(op4.input[0], 'foo/scale')
self.assertEqual(op4.output[0], 'foo/scale_2')
self.assertEqual(op5.input[0], 'foo/sum')
self.assertEqual(op5.output[0], 'foo/sum_2')
self.assertEqual(op6.input[0], 'foo/branch')
self.assertEqual(op6.input[1], 'test/branch_2')
self.assertEqual(op6.output[0], 'foo/branch_3')
self.assertEqual(op6.output[1], 'test/branch4')
