python - 如何在 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 documentation，它变成了 tf.FrozenModel。拥有可训练模型的唯一方法是通过 tf.loadLayersModel，这需要将 Keras 模型转换为 tensorflow.js 而不是 tensorflow savedmodel。但是，我也无法将转换后的 tensorflow savedmodel 转换为 Keras。是否可以将 PyTorch 模型导出到 tensorflow.js 并且它仍然可以训练？

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

编辑:这里有更多的细节。我正在尝试将 efficientnet 从 Pytorch 转换为 Tensorflow。

这会将 PyTorch efficientnet(来自名为 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>

加载 LayersModel:

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