tensorflowjs：编写 tensorflowjs 自定义层时如何扩展 tf.layers.conv2d

Question

我需要编写一个 tensorflowjs 的自定义层来导入我的模型。在原始 python 代码中，这是通过扩展 Conv2D 实现的，但这在 tfjs 代码中似乎是不可能的。寻求帮助。

错误信息

我是 tensorflow 的新手，我想知道如何将这个自定义层转换为 tfjs 代码，这一直困扰着我很长时间，任何帮助将不胜感激。

这里是python自定义层代码。

from keras.utils import conv_utils
from keras import backend as K
from keras.engine import InputSpec
from keras.layers import Conv2D


class PConv2D(Conv2D):
    def __init__(self, *args, n_channels=3, mono=False, **kwargs):
        super().__init__(*args, **kwargs)
        self.input_spec = [InputSpec(ndim=4), InputSpec(ndim=4)]

    def build(self, input_shape):        
        """Adapted from original _Conv() layer of Keras        
        param input_shape: list of dimensions for [img, mask]
        """

        if self.data_format == 'channels_first':
            channel_axis = 1
        else:
            channel_axis = -1

        if input_shape[0][channel_axis] is None:
            raise ValueError('The channel dimension of the inputs should be defined. Found `None`.')

        self.input_dim = input_shape[0][channel_axis]

        # Image kernel
        kernel_shape = self.kernel_size + (self.input_dim, self.filters)
        self.kernel = self.add_weight(shape=kernel_shape,
                                      initializer=self.kernel_initializer,
                                      name='img_kernel',
                                      regularizer=self.kernel_regularizer,
                                      constraint=self.kernel_constraint)
        # Mask kernel
        self.kernel_mask = K.ones(shape=self.kernel_size + (self.input_dim, self.filters))

        if self.use_bias:
            self.bias = self.add_weight(shape=(self.filters,),
                                        initializer=self.bias_initializer,
                                        name='bias',
                                        regularizer=self.bias_regularizer,
                                        constraint=self.bias_constraint)
        else:
            self.bias = None
        self.built = True

    def call(self, inputs, mask=None):
        '''
        We will be using the Keras conv2d method, and essentially we have
        to do here is multiply the mask with the input X, before we apply the
        convolutions. For the mask itself, we apply convolutions with all weights
        set to 1.
        Subsequently, we set all mask values >0 to 1, and otherwise 0
        ''' 

        # Both image and mask must be supplied
        if type(inputs) is not list or len(inputs) != 2:
            raise Exception('PartialConvolution2D must be called on a list of two tensors [img, mask]. Instead got: ' + str(inputs))

        # Create normalization. Slight change here compared to paper, using mean mask value instead of sum
        normalization = K.mean(inputs[1], axis=[1,2], keepdims=True)
        normalization = K.repeat_elements(normalization, inputs[1].shape[1], axis=1)
        normalization = K.repeat_elements(normalization, inputs[1].shape[2], axis=2)

        # Apply convolutions to image
        img_output = K.conv2d(
            (inputs[0]*inputs[1]) / normalization, self.kernel, 
            strides=self.strides,
            padding=self.padding,
            data_format=self.data_format,
            dilation_rate=self.dilation_rate
        )

        # Apply convolutions to mask
        mask_output = K.conv2d(
            inputs[1], self.kernel_mask, 
            strides=self.strides,
            padding=self.padding,            
            data_format=self.data_format,
            dilation_rate=self.dilation_rate
        )

        # Where something happened, set 1, otherwise 0        
        mask_output = K.cast(K.greater(mask_output, 0), 'float32')

        # Apply bias only to the image (if chosen to do so)
        if self.use_bias:
            img_output = K.bias_add(
                img_output,
                self.bias,
                data_format=self.data_format)

        # Apply activations on the image
        if self.activation is not None:
            img_output = self.activation(img_output)

        return [img_output, mask_output]

    def compute_output_shape(self, input_shape):
        if self.data_format == 'channels_last':
            space = input_shape[0][1:-1]
            new_space = []
            for i in range(len(space)):
                new_dim = conv_utils.conv_output_length(
                    space[i],
                    self.kernel_size[i],
                    padding=self.padding,
                    stride=self.strides[i],
                    dilation=self.dilation_rate[i])
                new_space.append(new_dim)
            new_shape = (input_shape[0][0],) + tuple(new_space) + (self.filters,)
            return [new_shape, new_shape]
        if self.data_format == 'channels_first':
            space = input_shape[2:]
            new_space = []
            for i in range(len(space)):
                new_dim = conv_utils.conv_output_length(
                    space[i],
                    self.kernel_size[i],
                    padding=self.padding,
                    stride=self.strides[i],
                    dilation=self.dilation_rate[i])
                new_space.append(new_dim)
            new_shape = (input_shape[0], self.filters) + tuple(new_space)
            return [new_shape, new_shape]

这是我的 tfjs 代码

import * as tf from '@tensorflow/tfjs';
import {convOutputLength, deconvLength, normalizeArray} from '@tensorflow/tfjs-layers/dist/utils/conv_utils';
import { InputSpec } from '@tensorflow/tfjs-layers/dist/engine/topology';


class PConv2D extends tf.layers.Layer {
  constructor(config) {
    super({config});
    // TODO(bileschi): Can we point to documentation on masking here?
    this.supportsMasking = true;
    this.inputSpec = [new InputSpec({ ndim: 4 }),new InputSpec({ ndim: 4 })];
    console.log(this.activation)
    if (config != undefined)
    {
        console.log(config.filters);
        this.kernelSize = config.kernelSize;
        this.filters = config.filters;
        this.strides = config.strides;
        this.padding = config.padding;
        this.dataFormat = config.dataFormat;
        this.dilationRate = config.dilationRate;
        this.useBias = config.useBias;
        this.biasInitializer = config.biasInitializer;
        this.biasConstraint = config.biasConstraint;
        this.kernelInitializer = config.kernelInitializer;
        this.kernelRegularizer = config.kernelRegularizer;
        this.activation = config.activation;

    }
    console.log(config);
    //console.log(config.name);
    console.log(this.input_spec);
  }
  build(inputShape) {

    const channelAxis =
        this.dataFormat === 'channelsFirst' ? 1 : inputShape.length - 1;
    if (inputShape[channelAxis] == null) {
      throw new ValueError(
          'The channel dimension of the inputs should be defined. ' +
          'Found `None`.');
    }

    const inputDim = inputShape[0][channelAxis];
    const kernelShape = this.kernelSize.concat([inputDim,this.filters]);

    // Image kernel

    this.kernel = this.addWeight(
        'img_kernel', kernelShape, 'float32', this.kernelInitializer,
        this.kernelRegularizer, true, this.kernelConstraint);

    this.kernel_mask = tf.ones(kernelShape);

    if (this.useBias) {
      this.bias = this.addWeight(
          'bias', [this.filters], 'float32', this.biasInitializer,
          this.biasRegularizer, true, this.biasConstraint);
      };

    this.built = true;
    this.inputSpec = [new InputSpec({ ndim: 4 }),new InputSpec({ ndim: 4 })];
    //this.inputSpec = [{ndim: 4},{ndim: 4}];


    };

  computeOutputShape(inputShape) {
    const space = (this.dataFormat === 'channelsLast') ?
      inputShape[0].slice(1, inputShape.length - 1) :
      inputShape.slice(2);

    const newSpace = [];

    for (let i = 0; i < space.length; ++i) {
        const newDim = convOutputLength(
            space[i], this.kernelSize[i], this.padding, this.strides[i],
            typeof this.dilationRate === 'number' ? this.dilationRate :
                                                    this.dilationRate[i]);
        newSpace.push(newDim);
    }
    let outputShape = [inputShape[0]];

  if (this.dataFormat === 'channelsLast') {
        outputShape = outputShape[0].concat(newSpace);
        outputShape.push(this.filters);
      } else {
        outputShape.push(this.filters);
        outputShape = outputShape.concat(newSpace);
      }

    return [outputShape,outputShape];

  }

  call(inputs, kwargs) {


    return tf.tidy(() => {
      // if (inputs.length !== 2) {
      //     throw new ValueError(
      //         'PartialConvolution2D must be called on a list of two tensors [img, mask]');
      // }

      let normalization = tf.mean(inputs[1], [1,2],true)
      normalization = normalization.tile([1, inputs[0].shape[1], inputs[0].shape[2],1])


      let a = tf.mul(inputs[0],inputs[1])
      a = tf.mul(a,normalization)

      let img_outputs = tf.conv2d(
          a,this.filters, this.kernel.read(),
          this.strides, this.padding, this.dataFormat,this.dilationRate);

      let mask_outputs = tf.conv2d(
          inputs[1],this.filters, this.kernel.read(),
          this.strides, this.padding, this.dataFormat,this.dilationRate);

      mask_outputs = K.cast(K.greater(mask_output, 0), 'float32')

      if (this.activation != null) {
        img_outputs = this.activation.apply(img_outputs);
      }
      return [img_outputs, mask_outputs];

    });


  }

  /**
   * If a custom layer class is to support serialization, it must implement
   * the `className` static getter.
   */
  static get className() {
    return 'PConv2D';
  }
  getConfig() {
    const config = super.getConfig();
    //Object.assign(config, {kernelSize: this.kernelSize});
    console.log(config)
    return config;
  }
}
tf.serialization.registerClass(PConv2D);  // Needed for serialization.

export function pconv2d() {
  return new PConv2D();
}

var myTester = new pconv2d()

Answer 1

错误是由

引起的

var myTester = new pconv2d()

pconv2d 需要 kernelSize 和过滤器。作为参数传递的配置对象的其他属性是可选的

new pconv2d({kernelSize: number, filters: number}) 将解决此问题。

但就像在 python 层的情况下所做的那样，pconv2d 可以继承自 tf.layers.conv2d。

这是一个简单的例子来演示它

class pconv2d extends tf.layers.conv2d {
  static className = 'MyCustomLayer';

  constructor(config) {
    super(config);
    // add other initialization
  }
  // add computation of the layer
}
tf.serialization.registerClass(pconv2d);
a = tf.ones([16, 16, 16, 16])
b = new MyCustomLayer({kernelSize: 2, filters: 2}).apply(a)
b.print()