def jax2tf_to_tfjs(test_case: ModelTestCase):
"""Converts the given `module` using the TFjs converter."""
# the model must be converted with with_gradient set to True to be able to
# convert the saved model to TF.js, as "PreventGradient" is not supported.
tf_fn = jax2tf.convert(test_case.apply,
with_gradient=True,
enable_xla=False)
# Create tf.Variables for the parameters. If you want more useful variable
# names, you can use `tree.map_structure_with_path` from the `dm-tree`
# package.
param_vars = tf.nest.map_structure(
lambda param: tf.Variable(param, trainable=True), test_case.variables)
# This is the function that will be stored in the SavedModel. Note this only
# supports a single argument, but we'd like to be able to pass more
# arguments to a Module's __call__ function, so we pass them as a list that
# we expand when passing it to `tf_fn`.
@tf.function(autograph=False, jit_compile=False)
def tf_graph(inputs):
return tf_fn(param_vars, *inputs)
s_fn = tf_graph.get_concrete_function(
_get_signatures(test_case.input_specs))
signatures = {tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: s_fn}
wrapper = _ReusableSavedModelWrapper(tf_graph, param_vars)
saved_model_options = tf.saved_model.SaveOptions(
experimental_custom_gradients=True)
with TempDir() as saved_model_path, TempDir() as tfjs_model_path:
tf.saved_model.save(wrapper,
saved_model_path,
signatures=signatures,
options=saved_model_options)
tfjs_converter.convert([saved_model_path, tfjs_model_path])
def saved_model_to_tfjs(input_dir, save_dir):
"""Convert SavedModel to TFJS model."""
print(f'\nConverting to TFJS:\nInput:{input_dir}\nOutput:{save_dir}\n')
converter.convert([
'--input_format=tf_saved_model', '--signature_name=serving_default',
'--control_flow_v2=True', '--skip_op_check', '--quantize_float16=True',
'--experiments=True', input_dir, save_dir
])
print('TFJS Conversion Success!')
def convert_to_tfjs(self, graph_def_path, output_names):
try:
from tensorflowjs.converters import converter
except ImportError:
return None
tfjs_path = os.path.join(self.test_data_directory, self._testMethodName + "_tfjs")
try:
converter.convert([graph_def_path, tfjs_path, '--input_format', 'tf_frozen_model',
'--output_node_names', ','.join(output_names)])
except ValueError:
return None
model_path = os.path.join(tfjs_path, 'model.json')
if not os.path.exists(model_path):
return None
return model_path
def convert_to_tflite(self, graph_def, feed_dict, outputs):
if not feed_dict:
return None # Can't make TFlite model with no inputs
tf_reset_default_graph()
with tf_session() as sess:
tf.import_graph_def(graph_def, name='')
sess_inputs = [
sess.graph.get_tensor_by_name(k) for k in feed_dict.keys()
]
sess_outputs = [sess.graph.get_tensor_by_name(n) for n in outputs]
converter = tf_lite.TFLiteConverter.from_session(
sess, sess_inputs, sess_outputs)
#converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS, # enable TensorFlow Lite ops.
tf.lite.OpsSet.SELECT_TF_OPS, # enable TensorFlow flex ops.
]
from tensorflow.lite.python.convert import ConverterError
try:
tflite_model = converter.convert()
tflite_path = os.path.join(self.test_data_directory,
self._testMethodName + ".tflite")
dir_name = os.path.dirname(tflite_path)
if dir_name:
os.makedirs(dir_name, exist_ok=True)
with open(tflite_path, 'wb') as f:
f.write(tflite_model)
return tflite_path
except ConverterError:
return None
def jax2tf_to_tfjs(module: lib.ModuleToConvert):
"""Converts the given `module` using the TFjs converter."""
with TempDir() as saved_model_path, TempDir() as converted_model_path:
# the model must be converted with with_gradient set to True to be able to
# convert the saved model to TF.js, as "PreventGradient" is not supported
saved_model_lib.convert_and_save_model(
module.apply,
module.variables,
saved_model_path,
input_signatures=[
tf.TensorSpec(shape=module.input_shape,
dtype=module.dtype,
name='input')
],
with_gradient=True,
compile_model=False,
enable_xla=False)
tfjs_converter.convert([saved_model_path, converted_model_path])
def convert_to_tfjs(self, graph_def_path, output_names):
try:
from tensorflowjs.converters import converter
except ImportError:
self.logger.warning(
"Tensorflowjs.converters package imports failed.")
return None
tfjs_path = os.path.join(self.test_data_directory,
self._testMethodName + "_tfjs")
try:
converter.convert([
graph_def_path, tfjs_path, '--input_format', 'tf_frozen_model',
'--output_node_names', ','.join(output_names)
])
except ValueError:
self.logger.warning("Convert tensorflowjs graph failed.")
return None
model_path = os.path.join(tfjs_path, 'model.json')
if not os.path.exists(model_path):
self.logger.warning("Tensorflowjs model path %s is empty.",
model_path)
return None
return model_path
def run(dryrun):
print('Welcome to TensorFlow.js Converter.')
input_path = [{
'type':
'input',
'name':
common.INPUT_PATH,
'message':
'Please provide the path of model file or '
'the directory that contains model files. \n'
'If you are converting TFHub module please provide the URL.',
'filter':
os.path.expanduser,
'validate':
lambda path: 'Please enter a valid path' if not path else True
}]
input_params = PyInquirer.prompt(input_path, style=prompt_style)
detected_input_format, normalized_path = detect_input_format(
input_params[common.INPUT_PATH])
input_params[common.INPUT_PATH] = normalized_path
formats = [{
'type': 'list',
'name': common.INPUT_FORMAT,
'message': input_format_message(detected_input_format),
'choices': input_formats(detected_input_format)
}, {
'type':
'list',
'name':
common.OUTPUT_FORMAT,
'message':
'What is your output format?',
'choices':
available_output_formats,
'when':
lambda answers: value_in_list(answers, common.INPUT_FORMAT, (
common.KERAS_MODEL, common.TFJS_LAYERS_MODEL))
}]
format_params = PyInquirer.prompt(formats,
input_params,
style=prompt_style)
message = input_path_message(format_params)
questions = [{
'type':
'input',
'name':
common.INPUT_PATH,
'message':
message,
'filter':
expand_input_path,
'validate':
lambda value: validate_input_path(value, format_params[common.
INPUT_FORMAT]),
'when':
lambda answers: (not detected_input_format)
}, {
'type':
'list',
'name':
common.SAVED_MODEL_TAGS,
'choices':
available_tags,
'message':
'What is tags for the saved model?',
'when':
lambda answers:
(is_saved_model(answers[common.INPUT_FORMAT]) and
(common.OUTPUT_FORMAT not in format_params or format_params[
common.OUTPUT_FORMAT] == common.TFJS_GRAPH_MODEL))
}, {
'type':
'list',
'name':
common.SIGNATURE_NAME,
'message':
'What is signature name of the model?',
'choices':
available_signature_names,
'when':
lambda answers:
(is_saved_model(answers[common.INPUT_FORMAT]) and
(common.OUTPUT_FORMAT not in format_params or format_params[
common.OUTPUT_FORMAT] == common.TFJS_GRAPH_MODEL))
}, {
'type':
'list',
'name':
common.QUANTIZATION_BYTES,
'message':
'Do you want to compress the model? '
'(this will decrease the model precision.)',
'choices': [{
'name': 'No compression (Higher accuracy)',
'value': None
}, {
'name': '2x compression (Accuracy/size trade-off)',
'value': 2
}, {
'name': '4x compression (Smaller size)',
'value': 1
}]
}, {
'type':
'input',
'name':
common.WEIGHT_SHARD_SIZE_BYTES,
'message':
'Please enter shard size (in bytes) of the weight files?',
'default':
str(4 * 1024 * 1024),
'when':
lambda answers: value_in_list(answers, common.OUTPUT_FORMAT,
(common.TFJS_LAYERS_MODEL))
}, {
'type':
'confirm',
'name':
common.SPLIT_WEIGHTS_BY_LAYER,
'message':
'Do you want to split weights by layers?',
'default':
False,
'when':
lambda answers: value_in_list(answers, common.INPUT_FORMAT,
(common.TFJS_LAYERS_MODEL))
}, {
'type':
'confirm',
'name':
common.SKIP_OP_CHECK,
'message':
'Do you want to skip op validation? \n'
'This will allow conversion of unsupported ops, \n'
'you can implement them as custom ops in tfjs-converter.',
'default':
False,
'when':
lambda answers: value_in_list(answers, common.INPUT_FORMAT, (
common.TF_SAVED_MODEL, common.TF_HUB_MODEL))
}, {
'type':
'confirm',
'name':
common.STRIP_DEBUG_OPS,
'message':
'Do you want to strip debug ops? \n'
'This will improve model execution performance.',
'default':
True,
'when':
lambda answers: value_in_list(answers, common.INPUT_FORMAT, (
common.TF_SAVED_MODEL, common.TF_HUB_MODEL))
}]
params = PyInquirer.prompt(questions, format_params, style=prompt_style)
output_options = [{
'type': 'input',
'name': common.OUTPUT_PATH,
'message': 'Which directory do you want to save '
'the converted model in?',
'filter': lambda path: update_output_path(path, params),
'validate': lambda path: len(path) > 0
}, {
'type':
'confirm',
'message':
'The output already directory exists, '
'do you want to overwrite it?',
'name':
'overwrite_output_path',
'default':
False,
'when':
lambda ans: output_path_exists(ans[common.OUTPUT_PATH])
}]
while (common.OUTPUT_PATH not in params
or output_path_exists(params[common.OUTPUT_PATH])
and not params['overwrite_output_path']):
params = PyInquirer.prompt(output_options, params, style=prompt_style)
arguments = generate_arguments(params)
print('converter command generated:')
print('tensorflowjs_converter %s' % ' '.join(arguments))
print('\n\n')
log_file = os.path.join(tempfile.gettempdir(), 'converter_error.log')
if not dryrun:
try:
converter.convert(arguments)
print('\n\nFile(s) generated by conversion:')
print("Filename {0:25} Size(bytes)".format(''))
total_size = 0
output_path = params[common.OUTPUT_PATH]
if os.path.isfile(output_path):
output_path = os.path.dirname(output_path)
for basename in sorted(os.listdir(output_path)):
filename = os.path.join(output_path, basename)
size = os.path.getsize(filename)
print("{0:35} {1}".format(basename, size))
total_size += size
print("Total size:{0:24} {1}".format('', total_size))
except BaseException:
exc_type, exc_value, exc_traceback = sys.exc_info()
lines = traceback.format_exception(exc_type, exc_value,
exc_traceback)
with open(log_file, 'a') as writer:
writer.write(''.join(line for line in lines))
print('Conversion failed, please check error log file %s.' %
log_file)
def _convert_tfjs_model(saved_model_path: str, destination_path: str):
converter.convert([
CONVERTER_SAVED_MODEL_INPUT_FLAG, CONVERTER_SERVING_TAG_FLAG,
CONVERTER_DEFAULT_SIGNATURE_FLAG, saved_model_path, destination_path
])
def testTFJSPredictExtractorWithKerasModel(self, multi_model,
multi_output):
if not _TFJS_IMPORTED:
self.skipTest('This test requires TensorFlow JS.')
input1 = tf.keras.layers.Input(shape=(1, ), name='input1')
input2 = tf.keras.layers.Input(shape=(1, ), name='input2')
inputs = [input1, input2]
input_layer = tf.keras.layers.concatenate(inputs)
output_layers = {}
output_layers['output1'] = (tf.keras.layers.Dense(
1, activation=tf.nn.sigmoid, name='output1')(input_layer))
if multi_output:
output_layers['output2'] = (tf.keras.layers.Dense(
1, activation=tf.nn.sigmoid, name='output2')(input_layer))
model = tf.keras.models.Model(inputs, output_layers)
model.compile(optimizer=tf.keras.optimizers.Adam(lr=.001),
loss=tf.keras.losses.binary_crossentropy,
metrics=['accuracy'])
train_features = {'input1': [[0.0], [1.0]], 'input2': [[1.0], [0.0]]}
labels = {'output1': [[1], [0]]}
if multi_output:
labels['output2'] = [[1], [0]]
example_weights = {'output1': [1.0, 0.5]}
if multi_output:
example_weights['output2'] = [1.0, 0.5]
dataset = tf.data.Dataset.from_tensor_slices(
(train_features, labels, example_weights))
dataset = dataset.shuffle(buffer_size=1).repeat().batch(2)
model.fit(dataset, steps_per_epoch=1)
src_model_path = tempfile.mkdtemp()
model.save(src_model_path)
dst_model_path = tempfile.mkdtemp()
converter.convert([
'--input_format=tf_saved_model',
'--saved_model_tags=serve',
'--signature_name=serving_default',
src_model_path,
dst_model_path,
])
model_specs = [config_pb2.ModelSpec(name='model1', model_type='tf_js')]
if multi_model:
model_specs.append(
config_pb2.ModelSpec(name='model2', model_type='tf_js'))
eval_config = config_pb2.EvalConfig(model_specs=model_specs)
eval_shared_models = [
self.createTestEvalSharedModel(
model_name='model1',
eval_saved_model_path=dst_model_path,
model_type='tf_js')
]
if multi_model:
eval_shared_models.append(
self.createTestEvalSharedModel(
model_name='model2',
eval_saved_model_path=dst_model_path,
model_type='tf_js'))
schema = text_format.Parse(
"""
feature {
name: "input1"
type: FLOAT
}
feature {
name: "input2"
type: FLOAT
}
feature {
name: "non_model_feature"
type: INT
}
""", schema_pb2.Schema())
tfx_io = test_util.InMemoryTFExampleRecord(
schema=schema, raw_record_column_name=constants.ARROW_INPUT_COLUMN)
feature_extractor = features_extractor.FeaturesExtractor(eval_config)
predictor = tfjs_predict_extractor.TFJSPredictExtractor(
eval_config=eval_config, eval_shared_model=eval_shared_models)
examples = [
self._makeExample(input1=0.0, input2=1.0, non_model_feature=0),
self._makeExample(input1=1.0, input2=0.0, non_model_feature=1),
]
with beam.Pipeline() as pipeline:
# pylint: disable=no-value-for-parameter
result = (
pipeline
| 'Create' >> beam.Create(
[e.SerializeToString() for e in examples], reshuffle=False)
| 'BatchExamples' >> tfx_io.BeamSource(batch_size=2)
|
'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts()
| feature_extractor.stage_name >> feature_extractor.ptransform
| predictor.stage_name >> predictor.ptransform)
# pylint: enable=no-value-for-parameter
def check_result(got):
try:
self.assertLen(got, 1)
got = got[0]
self.assertIn(constants.PREDICTIONS_KEY, got)
self.assertLen(got[constants.PREDICTIONS_KEY], 2)
for item in got[constants.PREDICTIONS_KEY]:
if multi_model:
self.assertIn('model1', item)
self.assertIn('model2', item)
if multi_output:
self.assertIn('Identity', item['model1'])
self.assertIn('Identity_1', item['model1'])
elif multi_output:
self.assertIn('Identity', item)
self.assertIn('Identity_1', item)
except AssertionError as err:
raise util.BeamAssertException(err)
util.assert_that(result, check_result, label='result')
def testTFJSPredictExtractorWithSingleOutputModel(self, multi_model,
multi_output,
batch_examples,
batch_inputs):
input1 = tf.keras.layers.Input(shape=(1,), name='input1')
input2 = tf.keras.layers.Input(shape=(1,), name='input2')
inputs = [input1, input2]
input_layer = tf.keras.layers.concatenate(inputs)
output_layers = {}
output_layers['output1'] = (
tf.keras.layers.Dense(1, activation=tf.nn.sigmoid,
name='output1')(input_layer))
if multi_output:
output_layers['output2'] = (
tf.keras.layers.Dense(1, activation=tf.nn.sigmoid,
name='output2')(input_layer))
model = tf.keras.models.Model(inputs, output_layers)
model.compile(
optimizer=tf.keras.optimizers.Adam(lr=.001),
loss=tf.keras.losses.binary_crossentropy,
metrics=['accuracy'])
train_features = {'input1': [[0.0], [1.0]], 'input2': [[1.0], [0.0]]}
labels = {'output1': [[1], [0]]}
if multi_output:
labels['output2'] = [[1], [0]]
example_weights = {'output1': [1.0, 0.5]}
if multi_output:
example_weights['output2'] = [1.0, 0.5]
dataset = tf.data.Dataset.from_tensor_slices(
(train_features, labels, example_weights))
dataset = dataset.shuffle(buffer_size=1).repeat().batch(2)
model.fit(dataset, steps_per_epoch=1)
src_model_path = tempfile.mkdtemp()
model.save(src_model_path)
dst_model_path = tempfile.mkdtemp()
converter.convert([
'--input_format=tf_saved_model',
'--saved_model_tags=serve',
'--signature_name=serving_default',
src_model_path,
dst_model_path,
])
model_specs = [config.ModelSpec(name='model1', model_type='tf_js')]
if multi_model:
model_specs.append(config.ModelSpec(name='model2', model_type='tf_js'))
eval_config = config.EvalConfig(model_specs=model_specs)
eval_shared_models = [
self.createTestEvalSharedModel(
model_name='model1',
eval_saved_model_path=dst_model_path,
model_type='tf_js')
]
if multi_model:
eval_shared_models.append(
self.createTestEvalSharedModel(
model_name='model2',
eval_saved_model_path=dst_model_path,
model_type='tf_js'))
desired_batch_size = 2 if batch_examples else None
predictor = tfjs_predict_extractor.TFJSPredictExtractor(
eval_config=eval_config,
eval_shared_model=eval_shared_models,
desired_batch_size=desired_batch_size)
predict_features = [
{
'input1': np.array([0.0], dtype=np.float32),
'input2': np.array([1.0], dtype=np.float32),
'non_model_feature': np.array([0]), # should be ignored by model
},
{
'input1': np.array([1.0], dtype=np.float32),
'input2': np.array([0.0], dtype=np.float32),
'non_model_feature': np.array([1]), # should be ignored by model
}
]
if batch_inputs:
predict_features = [{k: np.expand_dims(v, 0)
for k, v in p.items()}
for p in predict_features]
with beam.Pipeline() as pipeline:
# pylint: disable=no-value-for-parameter
result = (
pipeline
| 'Create' >> beam.Create(predict_features)
| 'FeaturesToExtracts' >>
beam.Map(lambda x: {constants.FEATURES_KEY: x})
| predictor.stage_name >> predictor.ptransform)
# pylint: enable=no-value-for-parameter
def check_result(got):
try:
self.assertLen(got, 2)
# We can't verify the actual predictions, but we can verify the keys.
for item in got:
self.assertIn(constants.PREDICTIONS_KEY, item)
if multi_model:
self.assertIn('model1', item[constants.PREDICTIONS_KEY])
self.assertIn('model2', item[constants.PREDICTIONS_KEY])
if multi_output:
self.assertIn('Identity',
item[constants.PREDICTIONS_KEY]['model1'])
self.assertIn('Identity_1',
item[constants.PREDICTIONS_KEY]['model1'])
elif multi_output:
self.assertIn('Identity', item[constants.PREDICTIONS_KEY])
self.assertIn('Identity_1', item[constants.PREDICTIONS_KEY])
except AssertionError as err:
raise util.BeamAssertException(err)
util.assert_that(result, check_result, label='result')
