将 TensorFlowJS MobileNet + KNN 保存到 TFLite

问题描述

我已经使用 TensorFlowJS 在 MobileNet logits 结果之上训练了一个 KNN。

我想知道如何将 MobileNet + KNN 的结果导出到 TFLite 模型。

const knn = knnClassifier.create()
const net = await mobilenet.load()

const handleTrain = (imgEl,label) => {
  const image = tf.browser.fromPixels(imgEl);
  const activation = net.infer(image,true);
  knn.addExample(activation,label)
}

解决方法

1.保存模型

保存模型此示例将文件保存到本机文件系统，或者如果您需要将其保存在其他位置，请选中 the documentation。

await model.save('file:///path/to/my-model');

在这一步之后，您应该有一个 JSON 文件和一个二进制权重文件。

2.从 TensorFlow.js Layers 模型转换为 Saved Model 格式

tfjs_model.json 是您从上一步获得的 model.json 的路径，而 saved_model 是您要保存 SavedModel 格式的路径。
您可以从 here 中阅读有关使用 TensorflowJS 转换器的更多信息。

tensorflowjs_converter --input_format=tfjs_layers_model --output_format=keras_saved_model tfjs_model.json saved_model

3.从 SavedModel 格式转换为 TFLite 格式

从 SavedModel 格式转换为 TFLite 是推荐的方法，如 per the documentation。

import tensorflow as tf

# Convert the model
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir) # path to the SavedModel directory
tflite_model = converter.convert()

# Save the model.
with open('model.tflite','wb') as f:
  f.write(tflite_model)