如何在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。谢谢。

编辑：以下是更多细节。我正在尝试将来自Pytorch的efficientnet转换为Tensorflow。

这将把来自geffnet库的PyTorch efficientnet转换为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'

- Stanley Zheng

---

你需要使用TensorFlow.js吗？你是否接受另一种解决方案，以便您可以在JavaScript中训练模型？ - Justin Harris

这个问题是针对 TF.js 的，但肯定也对在 JavaScript 中使用 GPU 支持训练模型的其他方法感到好奇。 - Stanley

太好了。我们一直在寻找方法，以便在JavaScript中训练最初来自PyTorch的模型，以促进联邦学习。我们研究了您提到的事情，但遇到了相同的限制。我们从其他人那里汲取了一些想法，并最终将一个原型组合在一起，以从PyTorch模型生成ONNX梯度图，然后在JavaScript中使用该梯度图进行InferenceSession。我正在将所有想法整合成教程，这是一个预览 https://github.com/juharris/train-pytorch-in-js/tree/init - Justin Harris

很棒！看起来很不错。谢谢分享。 - Stanley

我使用ONNX Runtime Web成功地训练了“PyTorch模型”，用于处理MNIST数字。教程可以在https://github.com/juharris/train-pytorch-in-js找到。您可以在https://juharris.github.io/train-pytorch-in-js上在浏览器中运行它。 - Justin Harris

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) 

- yudhiesh

---

尝试了一下，谢谢。不幸的是，转换为tensorflow.js时，它引入了Lambda层，这些层无法转换为tensorflow.js。更具体地说，我正在尝试将efficientnet转换为tensorflow.js。我已经编辑了问题以添加更多细节。谢谢！ - Stanley

@StanleyZheng 我明白了，嗯，我知道这不是你想要的，但为什么不直接使用Tensorflow等效版本并将其转换为Tensorflowjs呢？ - yudhiesh

有点不幸的是，Efficientnet 的 tensorflow 版本 tf.keras.applications.efficientnet 使用了 Rescale 和 Normalize 预处理层，而这两个层在 tensorflow.js 中并没有实现。我想，除了实现这两个层之外，没有更容易的解决办法了，感谢 @yudhiesh 的帮助。 - Stanley