Tensorflow Custom Training

  Tensorflow를 사용하는 가장 큰 이유는 아마도 편리하게 Deep Learning 모델을 만들고 훈련시킬 수 있다는 점일 것입니다. 특히 2.0부터 Keras가 공식적으로 포함되면서 모델을 만들기가 더 쉬워졌습니다. 물론 그 전에도 Keras를 따로 사용하면 되기는 했지만요. Keras가 대부분의 인기 있는 아키텍처와 유틸리티를 지원하기 때문에 많은 부분을 직접 구현하지 않고, Keras의 구현을 가져다 쓸 수 있죠.

  하지만 가끔은 Keras에서 아직 지원하지 않는 모델을 사용하고 싶을 때가 있고, 훈련 과정을 좀 더 세밀하게 제어하고 싶을 때가 있습니다. 이때 Tensorflow의 Custom Training 기능을 사용할 수 있습니다. 오늘은 Keras 패키지 밖의 Tensorflow 기능들도 맛보기로 사용해보겠습니다.

1. Dataset 정의

  이번에는 실제 데이터를 사용하지 않고 가상의 데이터를 만들어보겠습니다. 1차 선형 데이터를 만들 건데요. y = W * x + b 형태에 노이즈를 더해서 데이터를 만들겠습니다.

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

  tf.random.normal 함수는 첫번째 인자인 shape에 맞는 random value를 만들어냅니다. 이때 normal distribution (mean=0, stddev=1)에서 추출을 하고, mean과 stddev는 함수의 인자로 변경할 수 있습니다. 이 예시에서는 n (데이터의 수) 개만큼 random value를 만들고 있습니다. targets = features * w + b 만 하면 정확한 값이 나오기 때문에 여기에 random으로 만든 noise를 더해줍니다.

  그럼 위와 같은 1차 선형이지만 noise가 더해진 가상의 데이터가 만들어집니다.

  우리는 W = 5.0, b = 1.0 인 데이터를 가상으로 만들어 보겠습니다. 훈련용 데이터(1,000개)와 평가용 데이터(100개)를 각각 만듭니다.

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)

2. Model 정의

  Custom Model을 정의하는 방법은 여러가지가 있겠지만 가장 편리한 방법은 tf.keras.Model을 상속 받아서 쓰는 것입니다. tf.keras.Model을 상속할 때 가장 중요한 메서드는 __init__과 call입니다. __init__에서는 모델 초기화를 해주고, call은 입력값을 받아서 결과를 계산합니다. DNN을 구현하고 있다면 call 메서드에서 fowrad propagation을 하면 됩니다.

  이번 예시에서는 일반적인 DNN과 차별화(?)를 위해서 1차 함수를 사용해보겠습니다. (데이터를 생성할 때 사용한 것과 같은 함수 형태입니다).

class MyModel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        
        self.W = tf.Variable(1.0)
        self.b = tf.Variable(0.0)        
        
    def call(self, x):
        return self.W * x + self.b     

  __init__ 메소드에서 부모 클래스의 __init__을 호출해주고, W와 b 값을 임의로 초기화해줍니다. 여기에서는 1.0, 0.0을 각각 사용했는데 실제 상황에서는 보통 random 값을 사용하겠죠?. tf.Variable은 Tensorflow가 추적하면서 값을 업데이트(훈련)해야 하는 데이터에 사용합니다. 즉, 우리가 훈련시킬 parameter들은 tf.Variable이 돼야 합니다. 

  call 메소드는 인자로 입력값(x)을 받습니다. 그리고 현재 W와 b를 이용해서 결과를 계산해서 반환합니다. 실제 모델링 계산이 일어나는 핵심 부분입니다.

3. 훈련

  이제 훈련 데이터를 사용해서 Parameter를 훈련시키는 부분이 남았습니다. 훈련에는 Gradient descent를 사용하겠습니다. 그렇다면 당연히 loss 함수에 대한 각 parameter의 partial derivative을 구하고 Backpropagation으로 Parameter를 업데이트해야겠죠? 살짝 미간이 찌푸려질 때쯤 희소식이 있습니다. Backpropagation을 구현할 때 가장 골치 아픈 partial derivative를 Tensorflow가 알아서 구해줍니다. 바로 tf.GradientTape이 이 일을 해줍니다.

  GradientTape을 with 문을 사용해서 context로 만들어주면, 그 안에서 일어나는 모든 계산을 GradientTape이 기억하고, 나중에 Gradient를 계산해줍니다. 정확히는 GradientTape이 watch하는 값들을 기억하는데, 모든 tf.Variable은 자동으로 watch가 됩니다. 코드를 통해서 보겠습니다.

def train(model, features, outputs, learning_rate, loss_function):        
    with tf.GradientTape() as tape:
        loss = loss_function(outputs, model(features))
        dW, db = tape.gradient(loss, [model.W, model.b])
        model.W.assign_sub(learning_rate * dW)
        model.b.assign_sub(learning_rate * db)

  먼저 현재 Parameter 상태에서 loss를 구합니다. 그리고 tf.GradientTape.gradient 함수를 통해서 각 Parameter의 partial derivative를 구합니다. tf.GradientTape.gradient의 첫번째 인자는 loss, 두 번째 있자는 partial derivative를 구할 parameter의 list입니다. 이렇게 아주 간단히 gradient를 구했습니다.

  이제 각 Parameter를 업데이트 룰에 따라 업데이트 해줍니다. tf.Variable의 assign_sub는 기존 값에서 인자를 뺀 값으로 현재 값을 업데이트합니다. 좀 더 익숙하게는 W := W - learning_rate * dW입니다.

 

  이 예시에서는 batch를 사용하지 않고 전체 데이터를 한번에 사용하고 있습니다. 즉 train 함수를 한번 호출하고 나면 1 epoch이 끝납니다.

이제 저 함수를 사용해서 여러 epoch을 돌릴 수 있는 loop를 만듭니다.

model = MyModel()

loss_function = tf.keras.losses.MeanSquaredError()

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

print('W: {:.4f}'.format(model.W.value().numpy()))
print('b: {:.4f}'.format(model.b.value().numpy()))
print('Test loss: {:.4f}'.format(loss_function(test_targets, model(test_features))))

  loss 함수로는 Mean Sqaured Error를 사용합니다. 고맙게도 Keras에서 이미 만들어 놨습니다.

  총 20 epoch을 돌리겠습니다. 한번 train 한 후 새 Parameter를 사용해서 loss를 계산합니다.

Epoches:   0 - loss: 11.9659
Epoches:   1 - loss: 8.0205
Epoches:   2 - loss: 5.5017
Epoches:   3 - loss: 3.8920
Epoches:   4 - loss: 2.8622
Epoches:   5 - loss: 2.2025
Epoches:   6 - loss: 1.7796
Epoches:   7 - loss: 1.5081
Epoches:   8 - loss: 1.3336
Epoches:   9 - loss: 1.2213
Epoches:  10 - loss: 1.1489
Epoches:  11 - loss: 1.1023
Epoches:  12 - loss: 1.0722
Epoches:  13 - loss: 1.0527
Epoches:  14 - loss: 1.0401
Epoches:  15 - loss: 1.0319
Epoches:  16 - loss: 1.0266
Epoches:  17 - loss: 1.0232
Epoches:  18 - loss: 1.0209
Epoches:  19 - loss: 1.0195
W: 4.9553
b: 1.0427
Test loss: 0.8833

  최종적으로 W는 4.9553, b는 1.0427이 됐습니다. 실제 값인 5.0, 1.0과 매우 유사하군요.

4. Model에서 다른 함수를 사용한다면?

  앞에서는 데이터를 생성하는 부분과 Model에서 같은 1차 함수를 사용했습니다. 하지만 실제 상황에서는 데이터의 분포를 알 수가 없죠. 그래서 이번에는 Model에서 1차 함수가 아닌 2차 함수를 사용해보겠습니다. 데이터 생성은 1차 함수 형태로 그대로 두고, Model과 Train 부분만 살짝 수정합니다.

class MyModel_Quadratic(tf.keras.Model):
    def __init__(self):
        super().__init__()
        
        self.W1 = tf.Variable(1.0)
        self.W2 = tf.Variable(1.0)        
        self.b = tf.Variable(0.0)        
        
    def call(self, x):
        return (self.W1 * x ** 2) + (self.W2 * x) + self.b             

def train_quadratic(model, features, outputs, learning_rate, loss_function):        
    with tf.GradientTape() as tape:
        loss = loss_function(outputs, model(features))
        dW1, dW2, db = tape.gradient(loss, [model.W1, model.W2, model.b])
        model.W1.assign_sub(learning_rate * dW1)
        model.W2.assign_sub(learning_rate * dW2)
        model.b.assign_sub(learning_rate * db)
                    
model = MyModel_Quadratic()

loss_function = tf.keras.losses.MeanSquaredError()

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

print('W1: {:.4f}'.format(model.W1.value().numpy()))
print('W2: {:.4f}'.format(model.W2.value().numpy()))
print('b: {:.4f}'.format(model.b.value().numpy()))
print('Test loss: {:.4f}'.format(loss_function(test_targets, model(test_features))))

Epoches:   0 - loss: 11.6979
Epoches:   1 - loss: 8.2448
Epoches:   2 - loss: 5.9271
Epoches:   3 - loss: 4.3668
Epoches:   4 - loss: 3.3132
Epoches:   5 - loss: 2.5991
Epoches:   6 - loss: 2.1132
Epoches:   7 - loss: 1.7811
Epoches:   8 - loss: 1.5530
Epoches:   9 - loss: 1.3955
Epoches:  10 - loss: 1.2861
Epoches:  11 - loss: 1.2096
Epoches:  12 - loss: 1.1558
Epoches:  13 - loss: 1.1177
Epoches:  14 - loss: 1.0906
Epoches:  15 - loss: 1.0711
Epoches:  16 - loss: 1.0570
Epoches:  17 - loss: 1.0467
Epoches:  18 - loss: 1.0392
Epoches:  19 - loss: 1.0336
W1: 0.0688
W2: 4.9304
b: 0.9052
Test loss: 0.9270

  2차 항의 계수인 W1의 값은 거의 0에 가까워지고, W2는 5.0, b는 1.0에 가까워진 것을 볼 수 있습니다.

5. 마무리

  모델의 복잡도에 따라 실제 구현은 더 복잡해지겠지만 Custom Training을 하는 큰 틀은 같습니다.

1. Model을 정의한다. 이 때 훈련시킬 Parameter는 tf.Variable으로 정의한다.
2. call() 메소드에서 실제 계산(foward propagation)을 수행한다.
3. tf.GradientTape을 이용해서 각 gradient를 계산하다. (필요하면 mini batch 형태로 수행)
4. 계산된 gradient를 이용해서 Parameter를 업데이트한다.
5. 3과 4를 필요한 epoch만큼 반복한다.