[tensorflow] Custom Lambda layer / Custom layer

(tensorflow v2.4.0)

 

tensorflow에서 Lambda layer를 사용하면 사용자가 정의한 코드로 layer를 실행할 수 있으며, Sequential API model 안에서 임의의 함수로 실행됩니다.

Lambda layer는 다음과 같은 형태로 정의하여 사용할 수 있습니다.

tf.keras.layers.Lambda(lambda x: tf.abs(x))

이 Lambda layer는 input을 절대값으로 바꿔주는 역할을 하는 layer가 되는 것이죠.

 

그렇다면 기본적인 MNIST model을 어떻게 Lambda layer로 구성하는지 살펴보도록 하겠습니다.

import tensorflow as tf
from tensorflow.keras.datasets import mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = x_train / 255.
x_test = x_test / 255.

model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28,28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])
model.fit(x_train, y_train, epochs=5)
model.evaluate(x_test, y_test)

기본적인 MNIST model입니다. (28, 28) 크기의 입력을 Flatten()으로 1차원벡터로 변경해주고, Dense layer를 2개 쌓아서 softmax output을 통해서 예측하고 있습니다. 이 Sequential model에서 첫번째 Dense layer의 activation을 Lambda layer로 구성해보도록 하겠습니다.

 

첫번째 Dense layer의 activation을 제거하고, 그 역할을 Lambda Layer로 변경합니다.

model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28,28)),
    tf.keras.layers.Dense(128),
    tf.keras.layers.Lambda(lambda x: tf.abs(x)),
    tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])
model.fit(x_train, y_train, epochs=5)
model.evaluate(x_test, y_test)

아시다시피 relu는 값이 음수라면 0, 양수라면 그대로 activate해주는 activation function입니다.

하지만, 현재 Lambda layer는 값을 절대값으로 변경해주도록 구현되어 있으므로, 음수의 units을 비활성화하지 않고, 양수로 바꿔서 활성화시켜주게 됩니다.

relu를 사용할 때와 유사한 결과를 얻을 수 있습니다.

 

그럼 이제는 실제 relu와 동일한 역할을 하는 Lambda layer를 구현해보도록 하겠습니다.

import tensorflow.keras.backend as K
def my_relu(x):
    return K.maximum(0.0, x)

model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28,28)),
    tf.keras.layers.Dense(128),
    tf.keras.layers.Lambda(my_relu),
    tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])
model.fit(x_train, y_train, epochs=5)
model.evaluate(x_test, y_test)

my_relu라는 relu의 연산을 하도록 정의하였고, Lambda layer를 통해서 my_relu 함수를 sequential layer로 적용하였습니다.

 

What is a Layer ?

tensorflow에는 위와 같은 layer들이 미리 정의되어 있습니다.

그렇다면 각 layer들은 어떤 구조를 갖고 있을까요?

기본적으로 layer는 클래스이며 layer의 state인 parameter(weight)와 신경망에서 layer의 목적(입력->출력)을 달성하기 위한 Computation(forward pass)로 구성되어 있습니다.

 

따라서 만약 Simple Dense layer라면 이 layer는 다음과 같은 구조를 갖고 있습니다.

Simple Dense layer

Custom layer 구현

Layer는 class이며, 우리는 이 class를 상속하여 custom layer를 정의할 수 있습니다.

class SimpleDense(tf.keras.layers.Layer):
    def __init__(self, units=32):
        super(SimpleDense, self).__init__()
        self.units = units
    
    def build(self, input_shape): # Create the state of the layer (weights)
        w_init = tf.random_normal_initializer()
        self.w = tf.Variable(name='kernel',
                             initial_value=w_init(shape=(input_shape[-1], self.units), dtype='float32'),
                             trainable=True)
        
        b_init = tf.zeros_initializer()
        self.b = tf.Variable(name='bias',
                             initial_value=b_init(shape=(self.units,), dtype='float32'),
                             trainable=True)
    
    def call(self, inputs): # Defines the computation from inputs to outputs
        return tf.matmul(inputs, self.w) + self.b

위 custom layer는 Dense layer를 custom layer로 구현한 것입니다. __init__ 함수를 통해서 Layer class의 속성들을 상속받고, build와 call 함수를 오버라이딩하여 layer를 구성하게 됩니다.

build 함수는 layer의 초기 상태를 정의하며, weight와 bias가 초기화됩니다. 그리고 call 함수를 통해서 forward computation이 수행됩니다.

 

위의 SimpleDense로 이루어진 Sequence model의 학습입니다.

import numpy as np
xs = np.array([-1.0, 0.0, 1.0, 2.0, 3.0, 4.0], dtype=float)
ys = np.array([-3.0, -1.0, 1.0, 3.0, 5.0, 7.0], dtype=float)

model = tf.keras.Sequential([SimpleDense(units=1)])
model.compile(optimizer='sgd', loss='mse')
model.fit(xs, ys, epochs=500, verbose=0)
print(model.predict([10.0]))

그리고 model의 variables 변수를 통해서 모델의 weight와 bias를 확인할 수 있습니다.

 

위에서 구현한 custom SimpleDense는 tensorflow에서 제공하는 Dense layer와 동일한 연산을 수행하기 때문에, Dense layer대신 사용이 가능합니다. 위에서 구현한 간단한 MNIST model은 이제 다음과 같이 구성할 수 있습니다.

model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28,28)),
    SimpleDense(128),
    tf.keras.layers.Lambda(my_relu),
    tf.keras.layers.Dense(10, activation='softmax')
])

 

Activating Custom Layers

위에서 custom layer는 구현했지만, 무엇가 부족한 것이 있습니다.

바로 activation의 유무입니다. 

tensorflow에서 제공하는 layer들은 layer 자체에서 activation을 선택해서 activation function을 출력에 적용할 수 있습니다.

 

당연히, custom layer에서도 custom layer에 activation 적용이 가능합니다.

class SimpleDense(tf.keras.layers.Layer):
    def __init__(self, units=32, activation=None):
        super(SimpleDense, self).__init__()
        self.units = units
        self.activation = tf.keras.activations.get(activation)
    
    def build(self, input_shape): # Create the state of the layer (weights)
        w_init = tf.random_normal_initializer()
        self.w = tf.Variable(name='kernel',
                             initial_value=w_init(shape=(input_shape[-1], self.units), dtype='float32'),
                             trainable=True)
        
        b_init = tf.zeros_initializer()
        self.b = tf.Variable(name='bias',
                             initial_value=b_init(shape=(self.units,), dtype='float32'),
                             trainable=True)
    
    def call(self, inputs): # Defines the computation from inputs to outputs
        return self.activation(tf.matmul(inputs, self.w) + self.b)

위에서 구현한 Simple layer에서 인자로 activation을 추가로 읽고 있습니다. 기본값은 None이기 때문에, activation을 설정하지 않는다면, 적용되지 않은 단순 내적의 결과가 출력됩니다.

하지만, activation에 'relu'나 'sigmoid'와 같은 activation name을 인자로 넘겨주게 되면, 5 line에서 tf.keras.activations.get 메소드를 통해서 activation function API를 내부 변수로 저장하고, forward 계산 마지막에 19 line에서 activation이 적용된 출력이 리턴됩니다.

 

최종적으로 MNIST model은 다음과 같이 구성할 수 있습니다.

model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28,28)),
    SimpleDense(128, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

 

Quadratic layer 구현

이제 위에서 알아본 custom layer 구현을 사용해서 SimpleQuadratic layer를 구성해보도록 하겠습니다.

quadratic layer는 input \(x\)를 통해 output \(ax^2 + bx + c\)를 activation을 적용하여 출력합니다.

class SimpleQuadratic(Layer):
    def __init__(self, units=32, activation=None):
        '''Initializes the class and sets up the internal variables'''
        # YOUR CODE HERE
        super(SimpleQuadratic, self).__init__()
        self.units = units
        self.activation = tf.keras.activations.get(activation)
  
    def build(self, input_shape):
        '''Create the state of the layer (weights)'''
        # a and b should be initialized with random normal, c (or the bias) with zeros.
        # remember to set these as trainable.
        # YOUR CODE HERE
        a_init = tf.random_normal_initializer()
        self.a = tf.Variable(name='a',
                            initial_value=a_init(shape=(input_shape[-1], self.units), dtype=float),
                            trainable=True)
        b_init = tf.random_normal_initializer()
        self.b = tf.Variable(name='b',
                            initial_value=b_init(shape=(input_shape[-1], self.units), dtype=float),
                            trainable=True)
        c_init = tf.zeros_initializer()
        self.c = tf.Variable(name='c',
                            initial_value=c_init(shape=(self.units,), dtype=float),
                            trainable=True)
 
    def call(self, inputs):
        '''Defines the computation from inputs to outputs'''
        # YOUR CODE HERE
        x_squared = tf.math.square(inputs)
        x_squared_times_a = tf.matmul(x_squared, self.a)
        x_times_b = tf.matmul(inputs, self.b)
        x2a_plus_xb_plus_c = x_squared_times_a + x_times_b + self.c
        return self.activation(x2a_plus_xb_plus_c)

 

구현한 레이어로 MNIST model을 구성하여 학습시켜보도록 하겠습니다.

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  SimpleQuadratic(128, activation='relu'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

model.fit(x_train, y_train, epochs=5)
model.evaluate(x_test, y_test)

결과는 SimpleDense와 큰 차이는 나지 않습니다.

 

 

 

- 참조

Coursera - Custom Models, Layers, and Loss Functions with TensorFlow : week 3