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

【问题标题】：TensorflowJS: how to reset input/output shapes for pretrained model in TFJSTensorflowJS：如何在 TFJS 中重置预训练模型的输入/输出形状
【发布时间】：2021-03-30 12:12:05
【问题描述】：

对于pre-trained model in python，我们可以重置输入/输出形状：

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)

The source code

我正在尝试在 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)

 ...

我自己写的：

import * as tf from '@tensorflow/tfjs';

 export class DepthToSpace extends tf.layers.Layer {
    constructor() {
        super({});
    }

    computeOutputShape(shape: Array<number>) {
        // I think the issue is here
        // because the error occurs during initialization of the model
        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';
    }
}

使用模型：

 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，但形状为 1,128,128,32。

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

【问题讨论】：

标签： tensorflow tensorflow.js tensorflowjs-converter

【解决方案1】：

要从先前模型的层定义新模型，您需要使用tf.model

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

【讨论】：

谢谢，伙计！但我仍然有同样的错误：错误：输入 0 与层 p_re_lu 不兼容：输入形状的预期轴 1 具有值 96 但形状为 1,128,128,32。

【解决方案2】：

我尝试调试这个类：

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';
    }
}

并看到：当我不尝试重写大小时，computeOutputShape 方法只工作两次，当我尝试重置输入/输出时它工作 4 次。好吧，然后我打开模型的 JSON 文件并将输入从 [null, 96, 96, 32] 更改为 [null, 128, 128, 32] 并删除了这些行：

 // 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});

现在它适用于 128x128 图像。它看起来像上面的代码，添加层而不是重写它们。

【讨论】：