tensorflow.js custom layer 적용하기

 크롬 확장프로그램에 Tensorflow.js(이하 tfjs) 를 사용하여 keras 모델을 적용하던 중 한가지 문제점을 직면하였다. Keras의 AveragePooling1D 레이어가 tensorflow 2.6 버전을 기준으로 Masking을 지원하지만,  tfjs 3.9 버전에서는 Masking을 지원하지않는다. 따라서 Masking 을 지원하는 MeanPool 레이어를 구현하여야한다.

Custom Layer

class MyDenseLayer(tf.keras.layers.Layer):
  def __init__(self, num_outputs):
    super(MyDenseLayer, self).__init__()
    self.num_outputs = num_outputs

  def build(self, input_shape):
    self.kernel = self.add_weight("kernel",
                                  shape=[int(input_shape[-1]),
                                         self.num_outputs])

  def call(self, inputs):
    return tf.matmul(inputs, self.kernel)

 우선 Keras Docs(사용자 정의 층  |  TensorFlow Core)에서 설명하는 custom layer 는 위와 같다. 여기서 필수적으로 init 과 call 을 작성해주어야하며, 학습이 필요한 Layer의 경우 build 또한 작성해주어야한다. 우리가 작성하려는 MeanPool 레이어는 인풋을 받아 마스크가 있다면 마스킹을 한 뒤 평균을 구하는 레이어로 다음과 같이 작성할 수 있다.

class MeanPool(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        self.supports_masking = True
        super(MeanPool, self).__init__(**kwargs)

    def compute_mask(self, input, input_mask=None):
        # do not pass the mask to the next layers
        return None

    def call(self, x, mask=None):
        if mask is not None:
          mask = tf.cast(mask, x[0].dtype)
          mask = tf.expand_dims(mask, 2 )
          x *= mask
          return tf.reduce_sum(
              x, axis=1) / tf.reduce_sum(
              mask, axis=1)
        else:
          return tf.reduce_mean(x, axis=1)

    def compute_output_shape(self, input_shape):
          # remove temporal dimension
          return (input_shape[0], input_shape[2])

 해당 코드는 Keras에서 지원하는 AveragePooling1D 레이어를 적당히 수정한 것이다. 여기서 compute_output_shape는 인풋이 해당 레이어를 지나고 나서 아웃풋의 shape를 return 하는 method로 모델을 만들 때 내부적으로 계산에 쓰이니 필요하다. 하지만 이 레이어는 학습이 필요로 하지않으므로 build는 구현하지않았다. 현재 코드는 python으로 구현된 상태이니 이 것을 크롬 확장프로그램에서 사용하기 위해서는 javascript로 재작성 해주어야한다.

CustomLayer in Javascripts

class MeanPool extends tf.layers.Layer {
	constructor(config) {
		super(config);
		this.supportsMasking = true;
	}

    call(x, args){
      if (Object.keys(args).length > 0){

        return tf.tidy(() =>{
        var mask = tf.cast(args.mask, 'float32')
        mask = tf.expandDims(mask, 2)
        x = tf.mul(x[0], mask)

        return tf.div(tf.sum(x, 1), tf.sum(mask, 1))})
      }
      else
      	return tf.mean(x, 1)
    }


    computeOutputShape(input_shape){
      // remove temporal dimension
      return [input_shape[0], input_shape[2]]
    }

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

 우리는 Mask를 사용할 것이므로 Construct에서 supportMasking을 활성화해주어야한다.( 꼭 필요한 것인지는 모르겠으나 tfjs의 마스크를 사용하는 layer들은 이를 활성화시킨다.) 다른 부분들은 크게 다를 것이 없으나 한가지 알아야할 것은 tfjs 에서는 tensor 간의 연산에서 기호를 사용할 수 없다. 즉, a * b 같은 문법은 사용할 수 없으니 tf.mul 을 사용하자. 이렇게 custom Layer을 작성하였다면, 다음 구문을 꼭 추가하여서 이런 Layer가 있음을 알리자

tf.serialization.registerClass(MeanPool);

결과

const dummy = tf.tensor([[[1, 2, 3], [0, 0, 0], [0, 0, 0]], [[1, 1, 1], [0, 0, 0], [3, 3, 3]]]);
const mask = tf.tensor([[true, true, false], [true, false, true]])

const Meanlayer = new MeanPool();

console.log(Meanlayer.apply(dummy).dataSync());
console.log(Meanlayer.apply([dummy], {mask : mask}).dataSync());

 간단한 더미 코드를 작성하였고, 알아야할 것은 레이어를 테스트 하기 위해서는 apply method를 호출하면된다. 또한, console.log를 통해서 tensor 내부 값을 알 수 없으니 dataSync를 사용하도록하자. 결과는 다음과 같다.

[[0.3333, 0.6666, 1], [1.3333, 1.3333, 1.3333]]
[[0.5, 1, 1.5], [2, 2, 2]]

잘 작동하는 것을 확인해볼 수 있다.