[tensorflow] GradientTape의 활용법

TensorFlow 2.0의 자동 미분 기능인 GradientTape에 대하여 알아보겠습니다.

GradientTape은 자동 미분을 통해 동적으로 Gradient 값들을 확인해 볼 수 있다는 장점을 가지고 있습니다.

직접 print 해볼 수 있고, 또한 시각화 해 볼 수 있습니다.

아래 TensorFlow 공식 도큐먼트의 튜토리얼 진행을 통해 GradientTape의 자동미분 기능에 대하여 알아보도록 하겠습니다.

본 튜토리얼은 TensorFlow 공식 문서를 참조(reference) 하여 작성되었습니다.

출처

• TensorFlow - GradientTape
• TensorFlow - 전문가를 위한 빠른시작

import tensorflow as tf
import numpy as np

tf.__version__

'2.2.0'

샘플데이터 로드 (mnist)

mnist = tf.keras.datasets.mnist

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

x_train.min(), x_train.max()

(0.0, 1.0)

# 차원을 1 늘려줍니다
x_train = tf.expand_dims(x_train, -1)
x_test = tf.expand_dims(x_test, -1)

x_train.shape, x_test.shape

(TensorShape([60000, 28, 28, 1]), TensorShape([10000, 28, 28, 1]))

Dataset을 만듭니다

train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(10000).batch(32)
test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32)

STEP 1. Model을 정의합니다.

# layer 정의
input_ = tf.keras.layers.Input(shape=(28, 28, 1))
x = tf.keras.layers.Conv2D(32, 3, activation='relu')(input_)
x = tf.keras.layers.Conv2D(64, 3, activation='relu')(x)
x = tf.keras.layers.Flatten()(x)
x = tf.keras.layers.Dense(128, activation='relu')(x)
output_ = tf.keras.layers.Dense(10, activation='softmax')(x)

# model을 정의
model = tf.keras.models.Model(input_, output_)

model.summary()

Model: "model"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_1 (InputLayer)         [(None, 28, 28, 1)]       0         
_________________________________________________________________
conv2d_4 (Conv2D)            (None, 26, 26, 32)        320       
_________________________________________________________________
conv2d_5 (Conv2D)            (None, 24, 24, 64)        18496     
_________________________________________________________________
flatten_4 (Flatten)          (None, 36864)             0         
_________________________________________________________________
dense_8 (Dense)              (None, 128)               4718720   
_________________________________________________________________
dense_9 (Dense)              (None, 10)                1290      
=================================================================
Total params: 4,738,826
Trainable params: 4,738,826
Non-trainable params: 0
_________________________________________________________________

STEP 2. Loss Function을 정의합니다.

loss_function = tf.keras.losses.SparseCategoricalCrossentropy()

STEP 3. Optimizer를 정의합니다.

optimizer = tf.keras.optimizers.Adam()

STEP 4. Metric을 정의합니다.

train_loss = tf.keras.metrics.Mean()
train_acc = tf.keras.metrics.SparseCategoricalAccuracy()

test_loss = tf.keras.metrics.Mean()
test_acc = tf.keras.metrics.SparseCategoricalAccuracy()

STEP 5. Train/Test step 함수를 정의합니다.

GradientTape (그라디언트 테이프)

텐서플로는 자동 미분(주어진 입력 변수에 대한 연산의 그래디언트(gradient)를 계산하는 것) 을 위한 tf.GradientTape API를 제공합니다.

tf.GradientTape는 컨텍스트(context) 안에서 실행된 모든 연산을 테이프(tape)에 "기록"합니다.

그 다음 텐서플로는 후진 방식 자동 미분(reverse mode differentiation)을 사용해 테이프에 "기록된" 연산의 그래디언트를 계산합니다.

@tf.function
def train_step(images, labels):
    # 미분을 위한 GradientTape을 적용합니다.
    with tf.GradientTape() as tape:
        # 1. 예측 (prediction)
        predictions = model(images)
        # 2. Loss 계산
        loss = loss_function(labels, predictions)
    
    # 3. 그라디언트(gradients) 계산
    gradients = tape.gradient(loss, model.trainable_variables)
    
    # 4. 오차역전파(Backpropagation) - weight 업데이트
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    
    # loss와 accuracy를 업데이트 합니다.
    train_loss(loss)
    train_acc(labels, predictions)

@tf.function
def test_step(images, labels):
    # 1. 예측 (prediction)
    predictions = model(images)
    # 2. Loss 계산
    loss = loss_function(labels, predictions)
    
    # Test셋에 대해서는 gradient를 계산 및 backpropagation 하지 않습니다.
    
    # loss와 accuracy를 업데이트 합니다.
    test_loss(loss)
    test_acc(labels, predictions)

EPOCHS = 5

for epoch in range(EPOCHS):
    for images, labels in train_ds:
        train_step(images, labels)
        
    for test_images, test_labels in test_ds:
        test_step(test_images, test_labels)

    template = '에포크: {}, 손실: {:.5f}, 정확도: {:.2f}%, 테스트 손실: {:.5f}, 테스트 정확도: {:.2f}%'
    print (template.format(epoch+1,
                           train_loss.result(),
                           train_acc.result()*100,
                           test_loss.result(),
                           test_acc.result()*100))

에포크: 1, 손실: 0.01474, 정확도: 99.55%, 테스트 손실: 0.06000, 테스트 정확도: 98.74%
에포크: 2, 손실: 0.01392, 정확도: 99.57%, 테스트 손실: 0.06164, 테스트 정확도: 98.75%
에포크: 3, 손실: 0.01332, 정확도: 99.59%, 테스트 손실: 0.06439, 테스트 정확도: 98.75%
에포크: 4, 손실: 0.01273, 정확도: 99.61%, 테스트 손실: 0.06755, 테스트 정확도: 98.74%
에포크: 5, 손실: 0.01231, 정확도: 99.63%, 테스트 손실: 0.07177, 테스트 정확도: 98.72%