def testMultiTagSpec(self):
spec = native_module.create_module_spec(
batch_norm_module,
[({"training"}, {"training": True}),
({"inference"}, {"training": False})])
self.assertAllEqual(sorted(spec.get_tags()),
[set(["training"]), set(["inference"])])
def testGetExpectedImageSizeFromImageModuleInfo(self):
with tf.Graph().as_default():
spec = native_module.create_module_spec(image_module_fn_with_info)
self.assertAllEqual(image_util.get_expected_image_size(spec),
[2, 4])
m = module.Module(spec)
self.assertAllEqual(image_util.get_expected_image_size(m), [2, 4])
def testMultiSignatureSpec(self):
spec = native_module.create_module_spec(multi_signature_module)
self.assertAllEqual(sorted(spec.get_signature_names()),
["double", "square"])
self.assertAllEqual(list(spec.get_input_info_dict("double").keys()), ["x"])
self.assertAllEqual(list(spec.get_output_info_dict("double").keys()), ["y"])
self.assertAllEqual(list(spec.get_input_info_dict("square").keys()), ["z"])
self.assertAllEqual(list(spec.get_output_info_dict("square").keys()),
["z_out"])
def testNoSignaturesPresent(self):
def wrong_module_fn():
x = tf.placeholder(tf.float32, shape=[None, 3])
return tf.identity(x)
with self.assertRaises(ValueError) as cm:
spec = native_module.create_module_spec(wrong_module_fn)
self.assertIn("No signatures present", str(cm.exception))
def testModuleWithVariablesAndNoCheckpoint(self):
spec = native_module.create_module_spec(module_with_variables)
spec._create_impl(name="module", trainable=False, tags=None)
self.assertAllEqual(
[x.op.name for x in tf.global_variables()],
[
"module/weights",
"module/partition/part_0",
"module/partition/part_1",
"module/partition/part_2",
])
with tf.Session() as session:
session.run(tf.initializers.global_variables())
expected_values = [
[0.0, 0.0, 0.0],
[0.0, 0.0],
[0.0],
[0.0],
]
for a, b in zip(session.run(tf.global_variables()), expected_values):
self.assertAllEqual(a, b)
def testGetNumImageChannels(self):
with tf.Graph().as_default():
spec = native_module.create_module_spec(image_module_fn)
self.assertEqual(image_util.get_num_image_channels(spec), 3)
m = module.Module(spec)
self.assertEqual(image_util.get_num_image_channels(m), 3)
def testGetExpectedImageSizeFromShape(self):
spec = native_module.create_module_spec(image_module_fn)
self.assertAllEqual(image_util.get_expected_image_size(spec), [2, 4])
m = module.Module(spec)
self.assertAllEqual(image_util.get_expected_image_size(m), [2, 4])
def testDefaultTagSpec(self): spec = native_module.create_module_spec(multi_signature_module) self.assertAllEqual(sorted(spec.get_tags()), [set()])
def testGetNumImageChannels(self): spec = native_module.create_module_spec(image_module_fn) self.assertEqual(image_util.get_num_image_channels(spec), 3) m = module.Module(spec) self.assertEqual(image_util.get_num_image_channels(m), 3)
def testGetExpectedImageSizeFromShape(self): spec = native_module.create_module_spec(image_module_fn) self.assertAllEqual(image_util.get_expected_image_size(spec), [2, 4]) m = module.Module(spec) self.assertAllEqual(image_util.get_expected_image_size(m), [2, 4])