Блог

Я пытаюсь запустить преобразованную модель из репозитория: https://github.com/HasnainRaz/Fast-SRGAN . Что ж, преобразование прошло успешно. Но когда я попытался инициализировать модель, я увидел ошибку: "Unknown layer: TensorFlowOpLayer." . Если мы исследуем сохраненную модель, мы сможем увидеть TensorFlowOpLayer:

Структура модели

Как я понял, это и есть этот мир кода:

keras.layers.UpSampling2D(size=2, interpolation='bilinear')(layer_input) .

Я решил написать свой собственный класс «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 resu<
    }

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

Но это не работает. Может кто-нибудь помочь мне понять, как писать в метод «computeOutputShape»?
И второе недоразумение, почему на картинке выше мы видим следующий порядок слоев:

     Conv2D -> TensorFlowOpLayer -> PReLU
 

Как я понял, слой TensorFlowOpLayer — это «UpSampling2D» в коде python. Модель H5 была исследована на сайте: 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)
 

Инициализация модели в 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");
  }
 

как в коде 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)
 

Тогда как она используется:

     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;

    });
 

как в коде python:

 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]
 

Когда я добавил свой собственный класс «TensorFlowOpLayer», я вижу следующую ошибку:
"expected input1 to have shape [null,null,null,32] but got array with shape [1,96,96,3]."

Решена проблема. Проблема связана с версией кода и сохраненной моделью. Автор кода переработал код и не изменил сохраненную модель. Я переписал нужный класс:

 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 resu<
    }

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

и это работает.

Исходный код автора:

  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)