TensorflowJS：如何在TFJS中重置预训练模型的输入/输出形状

对于，我们可以重置输入/输出形状：

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)

我正在尝试在TFJS中执行同样的操作：

 // Load the model
 this.model = await tf.loadLayersModel('/assets/fast_srgan/model.json');

 // Define arbitrary spatial dims, and 3 channels.
 const inputs = tf.layers.input({shape: [null, null, 3]});
 
 // Trace out the graph using the input.
 const outputs = this.model.apply(inputs) as tf.SymbolicTensor;

 // Override the model.
 this.model = tf.model({inputs: inputs, outputs: outputs});

TFJS不支持模型中的某个层：

 ...

 u = keras.layers.Conv2D(filters, kernel_size=3, strides=1, padding='same')(layer_input)
 u = tf.nn.depth_to_space(u, 2) # <- TFJS does not support this layer
 u = keras.layers.PReLU(shared_axes=[1, 2])(u)

 ...

 tf.tidy(() => {
     let img = tf.browser.fromPixels(this.imgLr.nativeElement, 3);

     img = tf.div(img, 255);
     img = tf.expandDims(img, 0);

     let sr = this.model.predict(img) as tf.Tensor;

     sr = tf.mul(tf.div(tf.add(sr, 1), 2), 255).arraySync()[0];

     tf.browser.toPixels(sr as tf.Tensor3D, this.imgSrCanvas.nativeElement);
 });

但我得到了一个错误：

错误：输入0与层p_re_lu不兼容：输入形状的轴1应具有值96，但获得形状1128128,32

---

使用96x96像素图像对预训练模型进行训练。如果我使用96x96图像，它可以工作。但如果我尝试使用其他尺寸（例如128x128），它就不起作用。在python中，我们可以轻松地重置输入/输出形状。为什么它在JS中不起作用？

要从以前模型的层定义新模型，需要使用

tf.model

this.model = tf.model({inputs: inputs, outputs: outputs});

---

我尝试调试该类：

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 {
        const result = tf.depthToSpace(input[0], 2);
        return result;
    }

    static get className() {
        return 'TensorFlowOpLayer';
    }
}

---

现在它可以处理128x128个图像。它看起来像上面的代码，添加层而不是重写层。

谢谢，伙计！但我仍然有相同的错误：错误：输入0与层p_re_lu不兼容：输入形状的轴1的值应为96，但得到的形状为1128128,32。

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 {
        const result = tf.depthToSpace(input[0], 2);
        return result;
    }

    static get className() {
        return 'TensorFlowOpLayer';
    }
}

 // Define arbitrary spatial dims, and 3 channels.
 const inputs = tf.layers.input({shape: [null, null, 3]});
 
 // Trace out the graph using the input.
 const outputs = this.model.apply(inputs) as tf.SymbolicTensor;

 // Override the model.
 this.model = tf.model({inputs: inputs, outputs: outputs});