How to train a PyTorch model in TensorFlow.js?

Question

I'm looking to export my PyTorch model into tensorflow.js and have the ability to finetune it in tensorflow.js. To do this, I first convert PyTorch weights to ONNX, then to tensorflow, and finally use tensorflowjs_converter to convert to tensorflow.js. This results in an un-trainable model in TensorFlow.js. Is there any way to make this model trainable at one of these steps? The following is a minimal reproducible example.

First, defining a generic model and converting it in 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)

Next, I use onnx_tf to convert from ONNX to 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')

Finally, I use tensorflowjs_converter to convert to tensorflow.js with the command

tensorflowjs_converter --input_format=tf_saved_model model model_tfjs

However, when loading this in tensorflow.js with tf.loadGraphModel("model_tfjs/model.json"), it becomes a tf.FrozenModel according to tensorflow.js documentation. The only way to have a trainable model is by tf.loadLayersModel which requires a Keras model to be converted to tensorflow.js rather than a tensorflow savedmodel. However, I'm also unable to convert the converted tensorflow savedmodel into Keras. Is it possible to export a PyTorch model to tensorflow.js and have it still be trainable?

I've tried other libraries, pytorch2keras, onnx2keras, among others; they all seem to use lambda layers and therefore cannot be converted to tensorflow.js either. Thanks.

Edit: Here are additional details. I'm trying to convert an efficientnet from Pytorch to Tensorflow.

This converts the PyTorch efficientnet (from a library called geffnet) to ONNX. We can set either dynamic dimensions or static, but neither work.

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)

Next, we can convert to 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')

Finally we convert to tensorflow.js using tfjs converter.

tensorflowjs_converter --input_format=tensorflow_saved_model efficientnet_b0_tf efficientnet_b0_tfjs

A minimal test in tensorflow.js is as follows

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();

The preceeding example works as a tf.FrozenModel in inference, but cannot be trained. To be trained, a model must be converted to tensorflow.js from keras. My attempts to convert the python tensorflow model to keras have been unsuccessful. For example,

import tensorflow as tf

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

model.save(savepath)

This results in the following traceback:

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

Are you sure pytorch2keras does not work? You could try saving the model in the h5 format like the code below, which produces a tf.Model.

Saving the 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')

Converting the model:

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

Loading the LayersModel:

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