Skip to content

Commit

Permalink
Merge pull request hotosm#174 from hotosm/feature/tflite_support
Browse files Browse the repository at this point in the history
Support for .tflite
  • Loading branch information
kshitijrajsharma authored Oct 19, 2023
2 parents 62bb7f7 + b8163c2 commit e89c4fb
Showing 1 changed file with 11 additions and 0 deletions.
11 changes: 11 additions & 0 deletions backend/core/tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,17 @@ def train_model(
# Save the model in HDF5 format
model.save(os.path.join(output_path, "checkpoint.h5"))

logger.info(model.inputs)
logger.info(model.outputs)

# Convert the model to tflite for android/ios.
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

# Save the model.
with open(os.path.join(output_path, "checkpoint.tflite"), 'wb') as f:
f.write(tflite_model)

# now remove the ramp-data all our outputs are copied to our training workspace
shutil.rmtree(base_path)
training_instance.accuracy = float(final_accuracy)
Expand Down

0 comments on commit e89c4fb

Please sign in to comment.