如何在 TensorFlow.js 中训练 PyTorch 模型？

Question

我希望将我的 PyTorch 模型导出到 tensorflow.js 并能够在 tensorflow.js 中对其进行微调。 为此，我首先将 PyTorch 权重转换为 ONNX，然后转换为 tensorflow，最后使用tensorflowjs_converter转换为 tensorflow.js。 这会导致 TensorFlow.js 中的模型无法训练。 有什么方法可以让这个模型在这些步骤之一中训练？ 以下是一个最小的可重现示例。

首先，定义一个通用模型并在 PyTorch 中进行转换：

import torch
import torch.nn.functional as F


class ModelClass(torch.nn.Module):
    def __init__(self):
        super(ModelClass, self).__init__()
        self.fc1 = torch.nn.Linear(100, 10)
        self.fc2 = torch.nn.Linear(10, 1)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.sigmoid(self.fc2(x))
        return x


model = ModelClass()

example_input = torch.randn((1, 100), requires_grad=True)
print(model(example_input))

input_names = ["input0"]
output_names = ["output0"]
dynamic_axes = {'input0': {0: 'batch'}, 'output0': {0: 'batch'}}

torch_out = torch.onnx.export(
    model, example_input, 'model.onnx', export_params=True, verbose=True, input_names=input_names,
    output_names=output_names, dynamic_axes=dynamic_axes, opset_version=10,
    operator_export_type=torch.onnx.OperatorExportTypes.ONNX)

接下来，我使用onnx_tf将 ONNX 转换为 TensorFlow。

import onnx
import tensorflow as tf
from onnx_tf.backend import prepare

onnx_model = onnx.load('model.onnx')
tf_model = prepare(onnx_model)

tf_model.export_graph('model')

最后，我使用tensorflowjs_converter通过命令转换为 tensorflow.js

tensorflowjs_converter --input_format=tf_saved_model model model_tfjs

但是，当使用tf.loadGraphModel("model_tfjs/model.json")在 tensorflow.js 中加载它时，根据tensorflow.js 文档，它变成了tf.FrozenModel 。 拥有可训练模型的唯一方法是通过tf.loadLayersModel ，它需要将tf.loadLayersModel模型转换为 tensorflow.js 而不是 tensorflow 保存的模型。 但是，我也无法将转换后的 tensorflow 保存模型转换为 Keras。 是否可以将 PyTorch 模型导出到 tensorflow.js 并且仍然可以训练？

我尝试过其他库， pytorch2keras ， onnx2keras等； 它们似乎都使用 lambda 层，因此也无法转换为 tensorflow.js。 谢谢。

编辑：以下是其他详细信息。 我正在尝试将高效网络从 Pytorch 转换为 Tensorflow。

这会将 PyTorch 高效网络（来自名为 geffnet 的库）转换为 ONNX。 我们可以设置动态尺寸或静态尺寸，但都不起作用。

import onnx
import geffnet
import torch

efficientnet = 'efficientnet_b0'
DYNAMIC_SIZE = True

img_sizes = [224, 240, 260, 300, 380, 456, 528, 600, 672]
model_idx = int(efficientnet[-1]) # to find the correct static image size 

model = geffnet.create_model(
    efficientnet,
    in_chans=3,
    pretrained=True,
    exportable=True)

model.eval()

example_input = torch.randn((1, 3, img_sizes[model_idx], img_sizes[model_idx]), requires_grad=True)
model(example_input)

input_names = ["input0"]
output_names = ["output0"]
dynamic_axes = {'input0': {0: 'batch'}, 'output0': {0: 'batch'}}
if DYNAMIC_SIZE:
    dynamic_axes['input0'][2] = 'height'
    dynamic_axes['input0'][3] = 'width'

torch_out = torch.onnx.export(
    model, example_input, 'efficientnet_b0.onnx', export_params=True, verbose=False, input_names=input_names,
    output_names=output_names, dynamic_axes=dynamic_axes,
    opset_version=11, operator_export_type=torch.onnx.OperatorExportTypes.ONNX)

onnx_model = onnx.load('efficientnet_b0.onnx')
onnx.checker.check_model(onnx_model)

接下来，我们可以转换为 Tensorflow。

import onnx
from onnx_tf.backend import prepare

onnx_model = onnx.load(onnx_path)
tf_model = prepare(onnx_model)
tf_model.export_graph('efficientnet_b0_tf')

最后我们使用 tfjs 转换器转换为 tensorflow.js。

tensorflowjs_converter --input_format=tensorflow_saved_model efficientnet_b0_tf efficientnet_b0_tfjs

tensorflow.js 中的最小测试如下

const tf = require('@tensorflow/tfjs-node');

const getModel = async function () {
    const imgBase = await tf.loadGraphModel('file://./efficientnet_b0_tfjs/model.json');
    const x = tf.randomNormal([1, 224, 224, 3]);
    console.log(imgBase(x));
}
getModel();

前面的示例在推理中用作tf.FrozenModel ，但无法进行训练。 要进行训练，必须将模型从 keras 转换为 tensorflow.js。 我将 python tensorflow 模型转换为 keras 的尝试没有成功。 例如，

import tensorflow as tf

model = tf.keras.models.load_model('efficientnet_b0_tf')
print(model.summary())

model.save(savepath)

这导致以下回溯：

Traceback (most recent call last):
  File "graph2layers.py", line 29, in <module>
    graph2layers()
  File "graph2layers.py", line 18, in graph2layers
    print(model.summary())
AttributeError: '_UserObject' object has no attribute 'summary'

Answer 1

您确定pytorch2keras不起作用吗？ 您可以尝试以h5格式保存模型，如下面的代码，这会生成tf.Model 。

保存模型：

import torch
from pytorch2keras.converter import pytorch_to_keras

net = # your model
x = torch.randn(1, 3, 224, 224, requires_grad=False) # dummy input
k_model = pytorch_to_keras(net, x, [(3, None, None,)], verbose=True, names='short')
k_model.save('keras.h5')

转换模型：

tensorflowjs_converter --input_format keras \
                        <path-to-keras-model> \
                        <name-of-the-folder-to-save-js-model>

加载层模型：

const modelJson = require('<path-to-model.json>')
const model = await tf.loadLayersModel(modelJson)