Google Colabを用いてPytorchのモデルをTFLiteに変換

何の記事？

PytorchのモデルからTensorflow Liteへの変換を、Google colabを用いて行う方法です。

Pytorch → ONNX

import torch
import torchvision.models as models

model_name = "resnet50"

dummy_input = torch.randn((1, 3, 224, 224))
model = models.resnet50()
torch.onnx.export(model, dummy_input, "converted.onnx", verbose=True, input_names=['input'], output_names=['output'])

ONNX → Tensorflow

!pip install onnx_tf
from onnx_tf.backend import prepare
import onnx

TF_PATH = "./converted.pb" 
ONNX_PATH = "./converted.onnx" 
onnx_model = onnx.load(ONNX_PATH)  


tf_rep = prepare(onnx_model)  
tf_rep.export_graph(TF_PATH)

Tensorflow → TFLite

# Convert the model
converter = tf.lite.TFLiteConverter.from_saved_model(TF_PATH) # path to the SavedModel directory
converter.target_spec.supported_ops = [
  tf.lite.OpsSet.TFLITE_BUILTINS, # enable TensorFlow Lite ops.
  tf.lite.OpsSet.SELECT_TF_OPS # enable TensorFlow ops.
]
tflite_model = converter.convert()

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

最後に

細かい説明は省略します。コピペでOKです。