Tensorflow Custom Layer

Tensorflow Custom Training에서는 keras의 Model 클래스를 이용해서 Custom Model을 만들고 훈련시키는 법을 알아보았습니다.

이번에는 tf.keras.layers.Layer를 사용해서 Custom Layer를 만드는 법을 알아보겠습니다. Tensorflow Custom Training에서 만들었던 예시를 사용해서 이야기를 이어가보겠습니다.

tf.keras.Sequential을 쓰지 않고 직접 Model의 동작을 구성하기 위해서 Custom Model을 사용한다면, Custom Layer는 Model 내에서 특정 Layer의 동작을 직접 구성하기 위해서 사용합니다. 예를 들어 새로운 Dropout 알고리즘을 생각해냈다고 가정해보겠습니다. MyDropout이라는 Layer를 새로 만들고 나머지는 Tensorflow가 제공하는 것들을 그대로 사용할 수 있습니다. 물론 Custom Layer를 만들 정도의 상황이면 Model 자체도 손볼 것이 많기 때문에 Custom Model을 만들어 사용하는 경우가 많습니다.

tf.keras.layers.Layer

Custom Model이 tf.keras.Model 클래스를 상속받아 새로운 클래스를 만들었던 것과 비슷하게, Custom Layer는 tf.keras.layers.Layer 클래스를 상속받습니다. API도 매우 유사합니다.

• __init__(): Instance 생성 시에 호출됩니다.
• call(self, x): forward feeding 단계에서 호출됩니다. x 값을 이용해서 결과를 계산한 후 반환하면 됩니다.

Tensorflow Custom Training에서 만들었던 1차 함수 형태의 계산을 Custom Layer로 만들어보겠습니다.

class MyLinearLayer(tf.keras.layers.Layer):
    def __init__(self):
        super().__init__()

        initializer = tf.random_normal_initializer()

        self.W = tf.Variable(initializer(shape=(1,)))
        self.b = tf.Variable(initializer(shape=(1,)))

    def call(self, x):
        return self.W * x + self.b

Custom Model과 거의 내용이 같습니다. 지난번에는 초기값으로 고정값을 사용했는데 이번에는 tf.random_normal_initializer를 이용해서 Random value를 사용하겠습니다. 참고로 tf.random_normal_initializer의 기본값은 mean=0.0, stddev=0.05입니다.

이번에는 MyLinearLayer를 사용할 Custom Model입니다.

class MyModel(tf.keras.Model):
    def __init__(self):
        super().__init__()

        self.linear = MyLinearLayer()    


    def call(self, x):
        x = self.linear(x)

        return x

지난번에는 Model 내에서 parameter 관리와 prediction을 직접 했지만 이제는 그 부분을 MyLinearLayer로 옮기고 Model은 Layer의 구성에만 신경을 씁니다. 역할이 분리됐다는 점에서 더 나은 구성이라고 볼 수 있습니다.

apply_gradients()

훈련 부분은 한가지 중요한 차이를 제외하고는 지난번과 동일합니다. 지난번에는 각각의 gradient를 구한 후 update rule에 따라서 parameter를 직접 update했습니다. 이런 일이 워낙 흔하기 때문에 Tensorflow가 이 기능을 미리 구현해두었습니다.

def train(model, features, outputs, loss_function, optimizer):
    with tf.GradientTape() as tape:
        loss = loss_function(outputs, model(features))

    gradients = tape.gradient(loss, model.trainable_weights)        
    optimizer.apply_gradients(zip(gradients, model.trainable_weights))

• gradient를 구할 때 어떤 어떤 parameter에 대해서 gradient를 구하라고 수동으로 지정하지 않고, model.trainable_weights와 같이 모든 추적 중인 parameter를 지정해주십니다.
• parameter를 수동으로 update하지 않고, optimizer의 apply_gradients()를 이용해서 update합니다. apply_gradients()를 사용할 때는 (gradient, parameter)의 tuple을 전달하면 됩니다. 실제 업데이트는 각 optimizer가 정의한 방식에 따라 이루어집니다.

나머지 실행 부분

나머지 실행 부분도 지난번과 거의 같습니다. 수정한 train()을 호출하는 부분이 조금 수정되었습니다.

def generate_train_data(W, b, n):
    features = tf.random.normal([n])
    noise = tf.random.normal([n])
    targets = features * W + b + noise

    return features, targets

W = 5.0
b = 1.0
train_features, train_targets = generate_train_data(W, b, 1000)
test_features, test_targets = generate_train_data(W, b, 100)

model = MyModel()

loss_function = tf.keras.losses.MeanSquaredError()
optimizer = tf.keras.optimizers.Adam(learning_rate=0.1)

for epoch in range(100):
    train(model, train_features, train_targets, loss_function, optimizer)
    loss = loss_function(train_targets, model(train_features))
    print("Epoches: {:3d} - loss: {:.4f}".format(epoch, loss))

2차 함수 형태

지난번에 만들었던 2차 함수 형태도 한번 Custom Layer로 만들어 보겠습니다.

class MyQuadraticLayer(tf.keras.layers.Layer):
    def __init__(self):
        super().__init__()

        initializer = tf.random_normal_initializer()        

        self.W1 = tf.Variable(initializer(shape=(1,)))
        self.W2 = tf.Variable(initializer(shape=(1,)))
        self.b = tf.Variable(initializer(shape=(1,)))

    def call(self, x):
        return (self.W1 * x ** 2) + (self.W2 * x) + self.b   

class MyModel(tf.keras.Model):
    def __init__(self):
        super().__init__()

        self.quadratic = MyQuadraticLayer()    


    def call(self, x):
        x = self.quadratic(x)

        return x

거의 복붙 수준이네요.

일반화?

그런데 MyLinearLayer와 MyQuadraticLayer의 내용이 거의 비슷해서 찜찜합니다. Refactoring의 냄새가 나죠. 두 클래스를 하나로 합쳐보기로 합니다.

class MyPolynomialLayer(tf.keras.layers.Layer):
    def __init__(self, dimension):
        super().__init__()

        self.dimension = dimension

        initializer = tf.random_normal_initializer()        

        self.W = tf.Variable(initializer(shape=(dimension,)), name='weights')
        self.b = tf.Variable(initializer(shape=(1,)), name='bias')

    def call(self, x):        
        result = 0

        for i, w in enumerate(self.W.value()):
            result += (w * (x ** (self.dimension - i)))

        result = result + self.b

        return result  

class MyModel(tf.keras.Model):
    def __init__(self,):
        super().__init__()

        self.linear = MyPolynomialLayer(dimension=1)        

    def call(self, x):
        x = self.linear(x)

        return x

parameter를 W1, W2 로 하는 대신에 Tensor 형태인 W로 합쳤습니다. 그리고 디버깅의 편의를 위해서 W와 b에 이름도 추가해주었습니다. 이렇게 하지 않으면 Tensorflow log에 그냥 둘다 Variable:0이라고 나와서 구분이 잘 안 됩니다.

call() 내부가 조금 복잡해지기는 했는데, 주어진 차수에 따라서 값을 계산하는 내용이 들어갔습니다.

MyModel도 새 MyPolynomialLayer를 사용하도록 수정했습니다.

나름 Deep Netowrk?

MyPolynomialLayer는 일반적인 Keras Layer이기 때문에 연결해서 사용할 수도 있습니다. (왜인지는 모르지만) 1차 함수 형태로 계산 후에 이어서 2차 함수 형태로 계산을 하고 싶다고 가정해보겠습니다. 나름 Layer가 두개짜리 Deep Network(?)입니다.

class MyModel(tf.keras.Model):
    def __init__(self,):
        super().__init__()

        self.linear = MyPolynomialLayer(dimension=1)
        self.quadratic = MyPolynomialLayer(dimension=2)


    def call(self, x):
        x = self.linear(x)
        x = self.quadratic(x)

        return x

마무리

Tensorflow(Keras)가 제공하는 Layer들만 사용해도 충분히 좋은 모델을 만들 수 있습니다. Tensoflow가 업데이트될 때마다 새로운 Layer와 기능들을 제공하기도 하고요. 하지만 최신에 발표된 기술들을 실시간으로 포함하는 것은 현실적으로 어렵습니다. Custom Model과 Custom Layer 작성법을 알아두면 필요한 부분만 새로 작성하고 나머지는 Tensorflow가 제공하는 기능을 사용할 수 있기 때문에 비교적 간단하게 새로운 기술을 구현해볼 수 있습니다.