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

然而，当在 tensorflow.js 中使用 tf.loadGraphModel("model_tfjs/model.json") 加载它时，它根据 tensorflow.js documentation 变成了 tf.FrozenModel。拥有可训练模型的唯一方法是使用 tf.loadLayersModel，这需要将 Keras 模型转换为 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)

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

1 个答案: