如何将Pytorch模型转换为TFLite格式？

本文共计1055个文字，预计阅读时间需要5分钟。

目标是将服务器上用PyTorch训练的优秀模型转换为可在移动端运行的TFLite模型。最直接的想法是将PyTorch模型转换为TensorFlow模型，然后转换成TFLite。但目前这个转换目标尚未实现。

目标是想把在服务器上用pytorch训练好的模型转换为可以在移动端运行的tflite模型。

最直接的思路是想把pytorch模型转换为tensorflow的模型，然后转换为tflite。但是这个转换目前没有发现比较靠谱的方法。

经过调研发现最新的tflite已经支持直接从keras模型的转换，所以可以采用keras作为中间转换的桥梁，这样就能充分利用keras高层API的便利性。

转换的基本思想就是用pytorch中的各层网络的权重取出来后直接赋值给keras网络中的对应layer层的权重。

转换为Keras模型后，再通过tf.contrib.lite.TocoConverter把模型直接转为tflite.

下面是一个例子，假设转换的是一个两层的CNN网络。