Commit b0100ca2 authored by Martino Bertoni's avatar Martino Bertoni 🌋
Browse files

allowing exporting of uncompressed savedmodel

parent 62eca493
......@@ -6,7 +6,7 @@ import tensorflow_hub as hub
def export_smilespred(smilespred_path, destination,
tmp_path=None, clear_tmp=True):
tmp_path=None, clear_tmp=True, compress=True):
"""Export our Keras Smiles predictor to the TF-hub module format."""
from keras import backend as K
from chemicalchecker.tool.smilespred import Smilespred
......@@ -40,14 +40,14 @@ def export_smilespred(smilespred_path, destination,
signature_def_map={'serving_default': signature})
builder.save()
# now export savedmodel to module
export_savedmodel(tmp_path, destination)
export_savedmodel(tmp_path, destination, compress=compress)
# clean temporary folder
if clear_tmp:
shutil.rmtree(tmp_path)
def export_savedmodel(savedmodel_path, destination,
tmp_path=None, clear_tmp=True):
tmp_path=None, clear_tmp=True, compress=True):
"""Export Tensorflow SavedModel to the TF-hub module format."""
if tmp_path is None:
tmp_path = tempfile.mkdtemp()
......@@ -61,9 +61,12 @@ def export_savedmodel(savedmodel_path, destination,
sess.run(tf.global_variables_initializer())
module.export(tmp_path, sess)
print('DONE export_savedmodel')
# compress the exported files to destination
os.system("tar -cz -f %s --owner=0 --group=0 -C %s ." %
(destination, tmp_path))
if compress:
# compress the exported files to destination
os.system("tar -cz -f %s --owner=0 --group=0 -C %s ." %
(destination, tmp_path))
else:
shutil.copytree(tmp_path, destination)
# clean temporary folder
if clear_tmp:
shutil.rmtree(tmp_path)
......@@ -76,5 +79,6 @@ def export_batch(cc, destination_dir, datasets=None):
for ds in datasets:
s3 = cc.signature(ds, 'sign3')
pred_path = os.path.join(s3.model_path, 'smiles_final')
mdl_dest = os.path.join(destination_dir, ds[:2] + '.tar.gz')
export_smilespred(pred_path, mdl_dest)
mdl_dest = os.path.join(destination_dir, ds[:2])
export_smilespred(pred_path, mdl_dest, compress=False)
export_smilespred(pred_path, mdl_dest + '.tar.gz')
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment