def test_save_no_dist_restore_dist(self, model_and_input, distribution):
"""Save a model without DS, and restore it with DS."""
self.skipTest('Loading model with DS is not supported yet')
saved_dir = os.path.join(self.get_temp_dir(),
'test_save_no_dist_restore_dist')
model, output_name = model_and_input.get_model()
x_train, y_train, x_predict = model_and_input.get_data()
batch_size = model_and_input.get_batch_size()
self._train_model(model, x_train, y_train, batch_size)
predict_dataset = self._get_predict_dataset(x_predict, batch_size)
result_before_save = model.predict(predict_dataset)
saved_model.save(model, saved_dir)
with distribution.scope():
predict_dataset = distribution.experimental_distribute_dataset(
predict_dataset)
actual_data = next(iter(predict_dataset))
result_after_save = self._load_and_run_model(
saved_dir, actual_data)
self.assertAllEqual(result_before_save, result_after_save[output_name])
def _gen_uninitialized_variable(base_dir):
"""Generates a saved model with an uninitialized variable."""
class SubModule(module.Module):
"""A module with an UninitializedVariable."""
def __init__(self):
self.uninitialized_variable = resource_variable_ops.UninitializedVariable(
name="uninitialized_variable", dtype=dtypes.int64)
class Module(module.Module):
"""A module with an UninitializedVariable."""
def __init__(self):
super(Module, self).__init__()
self.sub_module = SubModule()
self.initialized_variable = variables.Variable(
1.0, name="initialized_variable")
# An UninitializedVariable with the same name as the variable in the
# SubModule, but with a different type.
self.uninitialized_variable = resource_variable_ops.UninitializedVariable(
name="uninitialized_variable", dtype=dtypes.float32)
@def_function.function(
input_signature=[tensor_spec.TensorSpec((), dtypes.float32)])
def compute(self, value):
return self.initialized_variable + value
to_save = Module()
saved_model.save(to_save,
export_dir=os.path.join(base_dir,
"UninitializedVariable"))
def testPostTrainingQuantize16x8(self):
"""Test for post-training quantization mode: activations/weights - int16/int8."""
saved_model_dir = os.path.join(self.get_temp_dir(), 'simple_savedmodel')
input_size = [5, 5, 3]
kernel_size = [3, 3, 1]
layer_name = 'test_conv2d'
input_0 = keras.layers.Input(shape=input_size)
layer_0 = keras.layers.Conv2D(
filters=kernel_size[-1],
kernel_size=kernel_size[0:2],
use_bias=False,
name=layer_name)(
input_0)
model = keras.models.Model(inputs=[input_0], outputs=[layer_0])
keras_layer = [layer for layer in model.layers if layer.name == layer_name
][0]
keras_layer.set_weights([
np.random.rand(
input_size[-1],
kernel_size[0],
kernel_size[1],
kernel_size[2],
).astype(np.float32)
])
saved_model.save(model, saved_model_dir)
model_coverage.test_saved_model(
saved_model_dir,
post_training_quantize_16x8=True,
model_input_size=input_size)
def test_save_dist_restore_no_dist(self, model_and_input, distribution):
"""Save a model with DS, and restore it without DS."""
self.skipTest('Saving model with DS is not supported yet')
saved_dir = os.path.join(self.get_temp_dir(),
'test_save_no_dist_restore_dist')
saved_dir_in_scope = os.path.join(saved_dir, _IN_SCOPE_SAVE_DIR)
saved_dir_out_of_scope = os.path.join(saved_dir,
_OUT_OF_SCOPE_SAVE_DIR)
with distribution.scope():
model, output_name = model_and_input.get_model()
x_train, y_train, x_predict = model_and_input.get_data()
batch_size = model_and_input.get_batch_size()
self._train_model(model, x_train, y_train, batch_size)
predict_dataset = self._get_predict_dataset(x_predict, batch_size)
result_before_save = model.predict(predict_dataset)
# save the model both in and out of the DS scope
saved_model.save(model, saved_dir_in_scope)
saved_model.save(model, saved_dir_out_of_scope)
actual_data = next(iter(predict_dataset))
result_load_from_save_in_scope = self._load_and_run_model(
saved_dir_in_scope, actual_data)
result_load_from_save_out_of_scope = self._load_and_run_model(
saved_dir_out_of_scope, actual_data)
self.assertAllEqual(result_before_save,
result_load_from_save_in_scope[output_name])
self.assertAllEqual(result_before_save,
result_load_from_save_out_of_scope[output_name])
def testKerasLenet(self):
"""Check that the output of PoplarExecutableRunner produces the same output as the original Graph execution.
"""
if utils.running_on_ipu_model():
self.skipTest(
"PoplarExecutableRunner only works with physical IPUs")
with tempfile.TemporaryDirectory() as tmp:
poplar_binaries_folder = os.path.join(tmp, "poplar")
model_path = os.path.join(tmp, "model")
weights_file = os.path.join(tmp, "weights.bin")
output_path = os.path.join(tmp, "output")
input_values = np.random.uniform(size=(1, 32, 32, 1))
input_file = "%s/input.bin" % tmp
with self.session() as sess:
self.configureIPU(poplar_binaries_folder, False)
with ops.device("/device:IPU:0"):
out, inp, model = instantiate_lenet()
utils.move_variable_initialization_to_cpu()
sess.run(global_variables_initializer())
utils.export_inputs_to_file([inp], input_file,
{inp: input_values})
# Run the model once to generate the poplar binaries.
reference_values = sess.run(out, {inp: input_values})
# Export the model & weights.
saved_model.save(model, model_path)
metadata_file = self.getSingleFileWithExt(poplar_binaries_folder,
"json")
executable_file = self.getSingleFileWithExt(
poplar_binaries_folder, "ipu_bin")
self.runPythonCommand(
(("./tensorflow/compiler/plugin/poplar/tools/"
"tensorflow_weights_extractor.py -o %s -s %s -m %s") %
(weights_file, model_path, metadata_file)).split())
self.runCommand((("./third_party/ipus/tools/PoplarExecutableRunner"
" --binaries %s,%s,%s "
"--output_folder=%s --strict") % (
executable_file,
weights_file,
input_file,
output_path,
)).split())
output_file = self.getSingleFileWithExt(output_path, "data")
with open(output_file, 'r') as f:
runner_values = np.array(json.load(f))
logging.info("Reference %s\nRunner: %s", reference_values,
runner_values)
self.assertAllClose(reference_values, runner_values)
def test_savedmodel(self):
class MyModule(module.Module):
@def_function.function(input_signature=[])
def foo(self):
return constant_op.constant([1])
saved_model.save(MyModule(), 'ram://my_module')
loaded = saved_model.load('ram://my_module')
self.assertAllEqual(loaded.foo(), [1])
def main(args):
if len(args) != 3:
print("Expected: {export_path} {ModuleName}")
print("Allowed ModuleNames:", MODULE_CTORS.keys())
return 1
_, export_path, module_name = args
module_ctor = MODULE_CTORS.get(module_name)
if not module_ctor:
print("Expected ModuleName to be one of:", MODULE_CTORS.keys())
return 2
os.makedirs(export_path)
tf_module = module_ctor()
options = save_options.SaveOptions(save_debug_info=True)
saved_model.save(tf_module, export_path, options=options)
def _gen_simple_while_loop(base_dir):
"""Generates a saved model with a while loop."""
class Module(module.Module):
"""A module with a while loop."""
@def_function.function(
input_signature=[tensor_spec.TensorSpec((), dtypes.float32)])
def compute(self, value):
acc, _ = control_flow_ops.while_loop(
cond=lambda acc, i: i > 0,
body=lambda acc, i: (acc + i, i - 1),
loop_vars=(constant_op.constant(0.0), value))
return acc
to_save = Module()
saved_model.save(to_save,
export_dir=os.path.join(base_dir, "SimpleWhileLoop"))
def test_save_load_io_device(self, model_and_input, distribution):
saved_dir = os.path.join(self.get_temp_dir(), 'io_device')
with distribution.scope():
model = model_and_input.get_model()
x_train, y_train, _ = model_and_input.get_data()
batch_size = model_and_input.get_batch_size()
self._train_model(model, x_train, y_train, batch_size)
call = model.__call__.get_concrete_function(tensor_spec.TensorSpec(None))
save_options = save_options_lib.SaveOptions(
experimental_io_device='/job:localhost')
saved_model.save(model, saved_dir, signatures=call, options=save_options)
load_options = load_options_lib.LoadOptions(
experimental_io_device='/job:localhost')
# Check that the model can be loaded and training continued without error.
with distribution.scope():
loaded_model = saved_model.load(saved_dir, options=load_options)
self._train_model(loaded_model, x_train, y_train, batch_size)
def test_save_dist_restore_dist(self, model_and_input, distribution_pair):
"""Save a model with DS, and restore it with potentially different DS."""
self.skipTest('Saving model with DS is not supported yet')
combinations.maybe_skip_test(self, distribution_pair.is_tpu_required,
distribution_pair.num_gpus_required)
saved_dir = os.path.join(self.get_temp_dir(),
'test_save_dist_restore_dist')
saved_dir_in_scope = os.path.join(saved_dir, _IN_SCOPE_SAVE_DIR)
saved_dir_out_of_scope = os.path.join(saved_dir,
_OUT_OF_SCOPE_SAVE_DIR)
dist_for_save = distribution_pair.strategy_1
dist_for_restore = distribution_pair.strategy_2
with dist_for_save.scope():
model, output_name = model_and_input.get_model()
x_train, y_train, x_predict = model_and_input.get_data()
batch_size = model_and_input.get_batch_size()
self._train_model(model, x_train, y_train, batch_size)
predict_dataset = self._get_predict_dataset(x_predict, batch_size)
result_before_save = model.predict(predict_dataset)
# save the model both in and out of the DS scope
saved_model.save(model, saved_dir_in_scope)
saved_model.save(model, saved_dir_out_of_scope)
with dist_for_restore.scope():
predict_dataset = dist_for_restore.experimental_distribute_dataset(
predict_dataset)
actual_data = next(iter(predict_dataset))
result_load_from_save_in_scope = self._load_and_run_model(
saved_dir_in_scope, actual_data)
result_load_from_save_out_of_scope = self._load_and_run_model(
saved_dir_out_of_scope, actual_data)
self.assertAllEqual(result_before_save,
result_load_from_save_in_scope[output_name])
self.assertAllEqual(result_before_save,
result_load_from_save_out_of_scope[output_name])
def main(args):
if len(args) != 3:
print("Expected: {export_path} {ModuleName}")
print("Allowed ModuleNames:", MODULE_CTORS.keys())
return 1
_, export_path, module_name = args
module_ctor, version = MODULE_CTORS.get(module_name)
if not module_ctor:
print("Expected ModuleName to be one of:", MODULE_CTORS.keys())
return 2
os.makedirs(export_path)
tf_module = module_ctor()
if version == 2:
options = save_options.SaveOptions(save_debug_info=True)
saved_model.save(tf_module, export_path, options=options)
else:
builder = saved_model.builder.SavedModelBuilder(export_path)
builder.add_meta_graph_and_variables(tf_module, ["serve"])
builder.save()
def testWeightsExportersNoMetadata(self):
""" Check that the weights extractor produces the same output with
TF v1 and v2 models."""
# Disable the IPU model
poplar_flags = os.environ.get("TF_POPLAR_FLAGS",
"").replace("--use_ipu_model", "")
with test.mock.patch.dict("os.environ",
{"TF_POPLAR_FLAGS": poplar_flags
}), tempfile.TemporaryDirectory() as tmp:
model_path_keras = os.path.join(tmp, "model_keras")
model_path_session = os.path.join(tmp, "model_session")
weights_keras = os.path.join(tmp, "weights_keras.bin")
weights_session = os.path.join(tmp, "weights_session.bin")
with self.session() as sess:
self.configureIPU()
with ops.device("/device:IPU:0"):
_, _, model = instantiate_lenet()
utils.move_variable_initialization_to_cpu()
sess.run(global_variables_initializer())
# Export the model & weights.
saved_model.save(model, model_path_keras)
Saver().save(sess, model_path_session)
self.runPythonCommand(
(("./tensorflow/compiler/plugin/poplar/tools/"
"tensorflow_weights_extractor.py -o %s -s %s") %
(weights_keras, model_path_keras)).split())
self.runPythonCommand(
(("./tensorflow/compiler/plugin/poplar/tools/"
"tensorflow_weights_extractor.py -o %s -s %s") %
(weights_session, model_path_session)).split())
with open(weights_session, 'rb') as s, open(weights_keras,
'rb') as k:
self.assertEqual(s.read(), k.read())
def _save_model(self, model, saved_dir): saved_model.save(model, saved_dir)
def _save_model(self, model, saved_dir): call = model.__call__.get_concrete_function(tensor_spec.TensorSpec(None)) saved_model.save(model, saved_dir, signatures=call)
def main(unused_argv):
model = get_model(MODEL_NAME.value)
path = os.path.join(TESTDATA_PATH.value, MODEL_NAME.value)
saved_model.save(model, path)
def save(self, path):
saved_model.save(self.model_, path)