def __init__(self,
hidden_1_dim=256,
hidden_2_dim=256,
input_dim=28 * 28,
classes=10):
super().__init__()
tf_utils.set_random_seed()
self.hidden_1_dim = hidden_1_dim
self.hidden_2_dim = hidden_2_dim
self.input_dim = input_dim
self.classes = classes
self.h1_weights = tf.Variable(
tf.random.normal([input_dim, hidden_1_dim]))
self.h2_weights = tf.Variable(
tf.random.normal([hidden_1_dim, hidden_2_dim]))
self.out_weights = tf.Variable(
tf.random.normal([hidden_2_dim, classes]))
self.h1_bias = tf.Variable(tf.random.normal([hidden_1_dim]))
self.h2_bias = tf.Variable(tf.random.normal([hidden_2_dim]))
self.out_bias = tf.Variable(tf.random.normal([classes]))
# Compile with dynamic batch dim.
self.predict = tf.function(
input_signature=[tf.TensorSpec([None, self.input_dim])])(
self.predict)
def concat2axis(self):
tf_utils.set_random_seed()
m = self.get_module()
a = tf.random.uniform([1, 5, 1], dtype=tf.float32)
b = tf.random.uniform([1, 5, 1], dtype=tf.float32)
dst = m.concat_zero_dim(a, b)
dst.assert_all_close()
def compare_backends(self, trace_function: Callable[[TracedModule], None],
modules: Modules) -> None:
"""Run the reference and target backends on trace_function and compare them.
Random seeds for tensorflow, numpy and python are set before each invocation
of trace_function.
Args:
trace_function: a function accepting a TracedModule as its argument.
"""
# Create Traces for each backend.
ref_trace = Trace(modules.ref_module, trace_function)
tar_traces = [
Trace(module, trace_function) for module in modules.tar_modules
]
# Run the traces through trace_function with their associated modules.
tf_utils.set_random_seed()
trace_function(TracedModule(modules.ref_module, ref_trace))
if FLAGS.log_all_traces:
logging.info(ref_trace)
for module, trace in zip(modules.tar_modules, tar_traces):
tf_utils.set_random_seed()
trace_function(TracedModule(module, trace))
if FLAGS.log_all_traces:
logging.info(trace)
# Compare each target trace of trace_function with the reference trace.
failed_backend_indices = []
error_messages = []
for i, tar_trace in enumerate(tar_traces):
logging.info("Comparing the reference backend '%s' with '%s'",
ref_trace.backend_id, tar_trace.backend_id)
traces_match, errors = Trace.compare_traces(ref_trace, tar_trace)
if not traces_match:
failed_backend_indices.append(i)
error_messages.extend(errors)
# Save the results to disk before validating.
ref_trace_dir = _get_trace_dir(modules.artifacts_dir, ref_trace)
ref_trace.save_plaintext(ref_trace_dir, FLAGS.summarize)
ref_trace.serialize(ref_trace_dir)
for tar_trace in tar_traces:
tar_trace_dir = _get_trace_dir(modules.artifacts_dir, tar_trace)
tar_trace.save_plaintext(tar_trace_dir, FLAGS.summarize)
tar_trace.serialize(tar_trace_dir)
# Validate results.
if failed_backend_indices:
# Extract info for logging.
failed_backends = [
tar_traces[i].backend_id for i in failed_backend_indices
]
error_list = ''.join(
[f'\n - {message}' for message in error_messages])
self.fail(
"Comparison between the reference backend and the following targets "
f"failed: {failed_backends}. Errors: {error_list}\n"
"See the logs above for more details about the non-matching calls."
)
def __init__(self):
super(SlimVisionModule, self).__init__()
tf_utils.set_random_seed()
model_path = posixpath.join(FLAGS.tf_hub_url, FLAGS.model, MODE)
hub_layer = hub.KerasLayer(model_path)
self.m = tf.keras.Sequential([hub_layer])
self.m.build(INPUT_SHAPE)
self.predict = tf.function(
input_signature=[tf.TensorSpec(INPUT_SHAPE)])(self.m.call)
def __init__(self):
super(LstmModule, self).__init__()
tf_utils.set_random_seed()
inputs = tf.keras.layers.Input(batch_size=None, shape=DYNAMIC_SHAPE[1:])
outputs = tf.keras.layers.LSTM(units=NUM_UNITS,
return_sequences=True)(inputs)
self.m = tf.keras.Model(inputs, outputs)
self.predict = tf.function(
input_signature=[tf.TensorSpec(DYNAMIC_SHAPE, tf.float32)])(self.m.call)
def __init__(self):
super().__init__()
tf_utils.set_random_seed()
model_path = posixpath.join(FLAGS.tf_hub_url, FLAGS.model, MODE)
hub_layer = hub.KerasLayer(model_path)
self.m = tf.keras.Sequential([hub_layer])
input_shape = get_input_shape()
self.m.build(input_shape)
self.predict = tf.function(input_signature=[tf.TensorSpec(input_shape)])(
self.m.call)
def compare_backends(self, trace_function: Callable[[TracedModule],
None]) -> None:
"""Run the reference and target backends on trace_function and compare them.
Random seeds for tensorflow, numpy and python are set before each invocation
of trace_function.
Args:
trace_function: a function accepting a TracedModule as its argument.
"""
# Create Traces for each backend.
ref_trace = Trace(_global_ref_module, trace_function)
tar_traces = [
Trace(module, trace_function) for module in _global_tar_modules
]
# Run the traces through trace_function with their associated modules.
tf_utils.set_random_seed()
trace_function(TracedModule(_global_ref_module, ref_trace))
if FLAGS.log_all_traces:
logging.info(ref_trace)
for module, trace in zip(_global_tar_modules, tar_traces):
tf_utils.set_random_seed()
trace_function(TracedModule(module, trace))
if FLAGS.log_all_traces:
logging.info(trace)
# Compare each target trace of trace_function with the reference trace.
failed_backend_indices = []
for i, tar_trace in enumerate(tar_traces):
logging.info("Comparing the reference backend '%s' with '%s'",
ref_trace.backend_name, tar_trace.backend_name)
traces_match = Trace.compare_traces(ref_trace, tar_trace)
if not traces_match:
failed_backend_indices.append(i)
# Save the results to disk before validating.
ref_trace_dir = _get_trace_dir(self._artifacts_dir, ref_trace)
ref_trace.save_plaintext(ref_trace_dir, FLAGS.summarize)
ref_trace.serialize(ref_trace_dir)
for tar_trace in tar_traces:
tar_trace_dir = _get_trace_dir(self._artifacts_dir, tar_trace)
tar_trace.save_plaintext(tar_trace_dir, FLAGS.summarize)
tar_trace.serialize(tar_trace_dir)
# Validate results.
if failed_backend_indices:
# Extract info for logging.
failed_backends = [
tar_traces[i].backend_name for i in failed_backend_indices
]
self.fail(
"Comparision between the reference backend and the following targets "
f"failed: {failed_backends}. The errors above show the inputs and "
"outputs of the non-matching calls.")
def lstm_module():
tf_utils.set_random_seed()
inputs = tf.keras.layers.Input(batch_size=NUM_BATCH, shape=INPUT_SHAPE[1:])
outputs = tf.keras.layers.LSTM(units=NUM_UNITS, return_sequences=True)(inputs)
model = tf.keras.Model(inputs, outputs)
module = tf.Module()
module.m = model
module.predict = tf.function(
input_signature=[tf.TensorSpec(INPUT_SHAPE, tf.float32)])(
model.call)
return module
def compare_backends(self, trace_function):
"""Run the reference and target backends on trace_function and compare them.
Random seeds for tensorflow, numpy and python are set before each invocation
of trace_function.
Args:
trace_function: a function accepting a TracedModule as its argument.
"""
# Create Traces for each backend.
ref_trace = Trace(self._ref_module, trace_function)
tar_traces = [Trace(module, trace_function) for module in self._tar_modules]
# Run the traces through trace_function with their associated modules.
tf_utils.set_random_seed()
trace_function(TracedModule(self._ref_module, ref_trace))
for module, trace in zip(self._tar_modules, tar_traces):
tf_utils.set_random_seed()
trace_function(TracedModule(module, trace))
# Compare each target trace of trace_function with the reference trace.
failed_backend_indices = []
for i, tar_trace in enumerate(tar_traces):
logging.info("Comparing the reference backend '%s' with '%s'",
ref_trace.backend, tar_trace.backend)
traces_match = Trace.compare_traces(ref_trace, tar_trace)
if not traces_match:
failed_backend_indices.append(i)
# Save the results to disk before validating.
ref_trace.save_plaintext(self._artifacts_dir, FLAGS.summarize)
for tar_trace in tar_traces:
tar_trace.save_plaintext(self._artifacts_dir, FLAGS.summarize)
# Validate results.
if failed_backend_indices:
# Extract info for logging.
failed_backends = [tar_traces[i].backend for i in failed_backend_indices]
failure_info = (
"Comparision between the reference backend and the following targets "
f"failed: {failed_backends}. The errors above show the inputs and "
"outputs the non-matching calls.")
# This condition is always True, but is useful for context in the logs.
self.assertEmpty(failed_backends, failure_info)
def CreateModule(input_dim=_FEATURE_SIZE, output_dim=1):
"""Creates a module for regression model training.
Args:
input_dim: input dimensionality
output_dim: output dimensionality
Returns:
model for linear regression
"""
tf_utils.set_random_seed()
# build a single layer model
inputs = tf.keras.layers.Input((input_dim))
outputs = tf.keras.layers.Dense(output_dim)(inputs)
model = tf.keras.Model(inputs, outputs)
return ModelTrain(model)
def initialize_model():
tf_utils.set_random_seed()
# Keras applications models receive input shapes without a batch dimension, as
# the batch size is dynamic by default. This selects just the image size.
input_shape = get_input_shape()[1:]
# If weights == 'imagenet', the model will load the appropriate weights from
# an external tf.keras URL.
weights = 'imagenet' if FLAGS.data == 'imagenet' else None
model = APP_MODELS[FLAGS.model](weights=weights,
include_top=FLAGS.include_top,
input_shape=input_shape)
if FLAGS.data == 'cifar10' and FLAGS.url:
model = load_cifar10_weights(model)
return model
def create_from_class(cls,
module_class: Type[tf.Module],
backend_info: "BackendInfo",
exported_names: Sequence[str] = (),
artifacts_dir: str = None):
"""Compile a tf.Module subclass to the target backend in backend_info.
Args:
module_class: The tf.Module subclass to compile.
backend_info: BackendInfo with the details for compiling module to IREE.
exported_names: Optional sequence representing the exported names to keep.
artifacts_dir: An optional string pointing to where compilation artifacts
should be saved. No compilation artifacts will be saved if this is not
provided.
"""
tf_utils.set_random_seed()
module_instance = module_class()
return cls.create_from_instance(module_instance, backend_info,
exported_names, artifacts_dir)
def create_from_class(cls,
module_class: Type[tf.Module],
backend_info: "BackendInfo",
exported_names: Sequence[str] = (),
artifacts_dir: str = None):
"""Compile a tf.Module subclass to the target backend in backend_info.
Args:
module_class: The tf.Module subclass to compile.
backend_info: BackendInfo with the details for compiling this module.
exported_names: Optional sequence representing the exported names to keep.
artifacts_dir: An optional string pointing to where compilation artifacts
should be saved. No compilation artifacts will be saved if this is not
provided.
"""
tf_utils.set_random_seed()
tflite_module_bytes = tf_module_to_tflite_module_bytes(
module_class, exported_names)
interpreters, compiled_paths = tflite_module_bytes_to_tflite_interpreters(
tflite_module_bytes, artifacts_dir)
module_name = module_class.__name__
return cls(module_name, backend_info, compiled_paths, interpreters)
def models():
tf.keras.backend.set_learning_phase(False)
tf_utils.set_random_seed()
input_shape = get_input_shape(FLAGS.data, FLAGS.model)
# keras model receives images size as input,
# where batch size is not specified - by default it is dynamic
if FLAGS.model in APP_MODELS:
weights = 'imagenet' if FLAGS.data == 'imagenet' else None
# if weights == 'imagenet' it will load weights from external tf.keras URL
model = APP_MODELS[FLAGS.model](weights=weights,
include_top=FLAGS.include_top,
input_shape=input_shape[1:])
if FLAGS.data == 'cifar10' and FLAGS.url:
file_name = 'cifar10' + FLAGS.model
# it will download model weights from publically available folder: PATH
# and save it to cache_dir=~/.keras and return path to it
weights_path = tf.keras.utils.get_file(
file_name,
os.path.join(
FLAGS.url,
'cifar10_include_top_{}_{}'.format(FLAGS.include_top,
FLAGS.model + '.h5')))
model.load_weights(weights_path)
else:
raise ValueError('Unsupported model', FLAGS.model)
module = tf.Module()
module.m = model
# specify input size with static batch size
# TODO(b/142948097): with support of dynamic shape
# replace input_shape by model.input_shape, so batch size will be dynamic (-1)
module.predict = tf.function(input_signature=[tf.TensorSpec(input_shape)])(
model.call)
return module
def initialize_model():
tf_utils.set_random_seed()
# Keras applications models receive input shapes without a batch dimension, as
# the batch size is dynamic by default. This selects just the image size.
input_shape = get_input_shape()[1:]
# If weights == 'imagenet', the model will load the appropriate weights from
# an external tf.keras URL.
weights = None
if FLAGS.use_external_weights and FLAGS.data == 'imagenet':
weights = 'imagenet'
model = APP_MODELS[FLAGS.model](weights=weights,
include_top=FLAGS.include_top,
input_shape=input_shape)
if FLAGS.use_external_weights and FLAGS.data == 'cifar10':
if not FLAGS.url:
raise ValueError(
'cifar10 weights cannot be loaded without the `--url` flag.')
model = load_cifar10_weights(model)
return model
def reinitialize(self):
"""Reinitializes all stateful variables."""
tf_utils.set_random_seed()
self._tf_module = self._constructor()