Tensorflow: converting H5 layer model to TFJS version leads to Unknown layer: TensorFlowOpLayer error when it works in TS - tensorflow

Question

I'm trying to run the converted model from the repository: https://github.com/HasnainRaz/Fast-SRGAN. Well, the conversion was successful. But when I tried to initialize the model, I saw the error: "Unknown layer: TensorFlowOpLayer.". If we will investigate the saved model, we can see TensorFlowOpLayer:
The model structure
As I understood it is this peace of code:
keras.layers.UpSampling2D(size=2, interpolation='bilinear')(layer_input).
I decided to write my own class "TensorFlowOpLayer".
import * as tf from '#tensorflow/tfjs';
export class TensorFlowOpLayer extends tf.layers.Layer {
constructor() {
super({});
}
computeOutputShape(shape: Array<number>) {
return [1, null, null, 32];
}
call(input_3): tf.Tensor {
const result = tf.layers.upSampling2d({ size: [2, 2], dataFormat: 'channelsLast', interpolation: 'bilinear' }).apply(input_3) as tf.Tensor;
return result;
}
static get className() {
return 'TensorFlowOpLayer';
}
}
But it doesn't work. Can someone help me to understand how to write to the method "computeOutputShape"?
And second misunderstanding, why on the picture above we see the next order of layers:
Conv2D -> TensorFlowOpLayer -> PReLU
As I understood the TensorFlowOpLayer layer is "UpSampling2D" in the python code. The H5 model was investigated through the site: https://netron.app
u = keras.layers.UpSampling2D(size=2, interpolation='bilinear')(layer_input)
u = keras.layers.Conv2D(self.gf, kernel_size=3, strides=1, padding='same')(u)
u = keras.layers.PReLU(shared_axes=[1, 2])(u)
The initializing of the model in TS:
async loadModel() {
this.model = await tf.loadLayersModel('/assets/fast_srgan/model.json');
const inputs = tf.layers.input({shape: [null, null, 32]});
const outputs = this.model.apply(inputs) as tf.SymbolicTensor;
this.model = tf.model({inputs: inputs, outputs: outputs});
console.log("Model has been loaded");
}
like in python code:
from tensorflow import keras
# Load the model
model = keras.models.load_model('models/generator.h5')
# Define arbitrary spatial dims, and 3 channels.
inputs = keras.Input((None, None, 3))
# Trace out the graph using the input:
outputs = model(inputs)
# Override the model:
model = keras.models.Model(inputs, outputs)
Then, how is it used:
tf.tidy(() => {
let img = tf.browser.fromPixels(this.imgLr.nativeElement, 3);
img = tf.div(img, 255.0);
img = tf.image.resizeNearestNeighbor(img, [96, 96]);
img = tf.expandDims(img, 0);
let sr = this.model.predict(img) as tf.Tensor;
});
like in python code:
def predict(img):
# Rescale to 0-1.
lr = tf.math.divide(img, 255)
# Get super resolution image
sr = model.predict(tf.expand_dims(lr, axis=0))
return sr[0]
When I added my own class "TensorFlowOpLayer" I see the next error:
"expected input1 to have shape [null,null,null,32] but got array with shape [1,96,96,3]."

Answer 1

Solved the issue. The issue related to the version of the code and the saved model. The author of the code refactored the code and didn't change the saved model. I rewrote the needed class:
import * as tf from '#tensorflow/tfjs';
export class DepthToSpace extends tf.layers.Layer {
constructor() {
super({});
}
computeOutputShape(shape: Array<number>) {
return [null, ...shape.slice(1, 3).map(x => x * 2), 32];
}
call(input): tf.Tensor {
input = input[0];
const result = tf.depthToSpace(input, 2);
return result;
}
static get className() {
return 'TensorFlowOpLayer';
}
}
and it works.
The author's original code is:
u = keras.layers.Conv2D(filters, kernel_size=3, strides=1, padding='same')(layer_input)
u = tf.nn.depth_to_space(u, 2)
u = keras.layers.PReLU(shared_axes=[1, 2])(u)