TensorFlow와 tf.GradientTape를 활용한 딥러닝

본 글은 Claude Sonnet 4로 작성 후, 수정한 내용입니다.

TensorFlow의 저수준 API인 tf.GradientTape를 활용한 딥러닝 학습 과정을 살펴보겠습니다. PyTorch에 익숙한 분들이라면 분명 Keras의 고수준 API가 답답했을 텐데, tf.GradientTape는 훨씬 직관적이고 세밀한 제어를 제공합니다.

tf.GradientTape란?

tf.GradientTape는 TensorFlow에서 자동미분을 위한 컨텍스트 관리자입니다. 테이프에 연산을 기록하고, 나중에 그 연산들의 gradient를 계산하는 역할을 합니다. PyTorch의 autograd와 매우 유사한 개념이죠.

import tensorflow as tf

# 간단한 예제
with tf.GradientTape() as tape:
    x = tf.Variable(3.0)
    y = x ** 2
    
# dy/dx 계산
dy_dx = tape.gradient(y, x)
print(dy_dx)  # 6.0

기본 학습 루프 구현

이제 실제 딥러닝 모델을 위한 학습 루프를 구현해보겠습니다.

import tensorflow as tf
import numpy as np

# 모델 정의
model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu', input_shape=(784,)),
    tf.keras.layers.Dense(32, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 손실 함수와 옵티마이저
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)

# 메트릭 정의
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
val_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()

def train_step(x, y):
    with tf.GradientTape() as tape:
        # 순전파
        predictions = model(x, training=True)
        loss = loss_fn(y, predictions)
    
    # 역전파
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    
    # 메트릭 업데이트
    train_accuracy.update_state(y, predictions)
    
    return loss

def val_step(x, y):
    predictions = model(x, training=False)
    val_accuracy.update_state(y, predictions)
    return loss_fn(y, predictions)

메모리 효율적인 대용량 데이터 처리

대용량 데이터를 처리할 때는 TFRecord를 활용하는 것이 좋습니다. 앞서 참고한 블로그의 방법을 응용해보겠습니다.

def create_dataset_from_tfrecord(file_pattern, batch_size=32):
    """TFRecord로부터 데이터셋 생성"""
    
    def parse_example(example_proto):
        features = {
            'image': tf.io.FixedLenFeature([], tf.string),
            'label': tf.io.FixedLenFeature([], tf.int64)
        }
        parsed = tf.io.parse_single_example(example_proto, features)
        
        # 이미지 디코딩
        image = tf.io.decode_raw(parsed['image'], tf.uint8)
        image = tf.cast(image, tf.float32) / 255.0
        image = tf.reshape(image, [28, 28, 1])
        
        return image, parsed['label']
    
    # 파일 로드
    filenames = tf.io.gfile.glob(file_pattern)
    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=tf.data.AUTOTUNE)
    
    # 전처리 파이프라인
    dataset = dataset.map(parse_example, num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.batch(batch_size)
    dataset = dataset.prefetch(tf.data.AUTOTUNE)
    
    return dataset

# 데이터셋 생성
train_dataset = create_dataset_from_tfrecord("train_*.tfrecord", batch_size=64)
val_dataset = create_dataset_from_tfrecord("val_*.tfrecord", batch_size=64)

전체 학습 루프

이제 모든 요소를 결합한 완전한 학습 루프를 구현해보겠습니다.

def train_model(model, train_dataset, val_dataset, epochs=10):
    """모델 학습 함수"""
    
    # 학습 히스토리 저장
    history = {'train_loss': [], 'train_acc': [], 'val_loss': [], 'val_acc': []}
    
    for epoch in range(epochs):
        print(f"Epoch {epoch + 1}/{epochs}")
        
        # 메트릭 초기화
        train_accuracy.reset_states()
        val_accuracy.reset_states()
        
        # 학습 단계
        train_loss_sum = 0
        train_batches = 0
        
        for x_batch, y_batch in train_dataset:
            loss = train_step(x_batch, y_batch)
            train_loss_sum += loss
            train_batches += 1
        
        # 검증 단계
        val_loss_sum = 0
        val_batches = 0
        
        for x_batch, y_batch in val_dataset:
            val_loss = val_step(x_batch, y_batch)
            val_loss_sum += val_loss
            val_batches += 1
        
        # 평균 계산
        avg_train_loss = train_loss_sum / train_batches
        avg_val_loss = val_loss_sum / val_batches
        
        # 결과 출력
        print(f"Train Loss: {avg_train_loss:.4f}, Train Acc: {train_accuracy.result():.4f}")
        print(f"Val Loss: {avg_val_loss:.4f}, Val Acc: {val_accuracy.result():.4f}")
        
        # 히스토리 저장
        history['train_loss'].append(float(avg_train_loss))
        history['train_acc'].append(float(train_accuracy.result()))
        history['val_loss'].append(float(avg_val_loss))
        history['val_acc'].append(float(val_accuracy.result()))
    
    return history

고급 활용: 그래디언트 클리핑과 정규화

실제 프로젝트에서는 그래디언트 클리핑이나 정규화를 적용해야 하는 경우가 많습니다.

def advanced_train_step(x, y, clip_norm=1.0):
    """고급 학습 단계 (그래디언트 클리핑 포함)"""
    
    with tf.GradientTape() as tape:
        predictions = model(x, training=True)
        
        # 기본 손실
        loss = loss_fn(y, predictions)
        
        # L2 정규화 추가
        l2_loss = tf.add_n([tf.nn.l2_loss(v) for v in model.trainable_variables])
        total_loss = loss + 0.01 * l2_loss
    
    # 그래디언트 계산
    gradients = tape.gradient(total_loss, model.trainable_variables)
    
    # 그래디언트 클리핑
    clipped_gradients, _ = tf.clip_by_global_norm(gradients, clip_norm)
    
    # 옵티마이저 적용
    optimizer.apply_gradients(zip(clipped_gradients, model.trainable_variables))
    
    train_accuracy.update_state(y, predictions)
    return loss

모델 평가와 예측

학습된 모델을 평가하고 예측하는 함수도 구현해보겠습니다.

def evaluate_model(model, test_dataset):
    """모델 평가"""
    
    test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
    test_loss = tf.keras.metrics.Mean()
    
    for x_batch, y_batch in test_dataset:
        predictions = model(x_batch, training=False)
        loss = loss_fn(y_batch, predictions)
        
        test_accuracy.update_state(y_batch, predictions)
        test_loss.update_state(loss)
    
    return test_loss.result().numpy(), test_accuracy.result().numpy()

def predict_batch(model, x_batch):
    """배치 예측"""
    predictions = model(x_batch, training=False)
    return tf.nn.softmax(predictions).numpy()

실제 사용 예제

# 모델 인스턴스 생성
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(28, 28, 1)),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Conv2D(64, 3, activation='relu'),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 데이터셋 준비
train_dataset = create_dataset_from_tfrecord("train_*.tfrecord", batch_size=64)
val_dataset = create_dataset_from_tfrecord("val_*.tfrecord", batch_size=64)

# 학습 실행
history = train_model(model, train_dataset, val_dataset, epochs=10)

# 모델 평가
test_dataset = create_dataset_from_tfrecord("test_*.tfrecord", batch_size=64)
test_loss, test_acc = evaluate_model(model, test_dataset)
print(f"Test Loss: {test_loss:.4f}, Test Accuracy: {test_acc:.4f}")

정리

tf.GradientTape를 활용하면 PyTorch와 유사한 방식으로 딥러닝 모델을 학습할 수 있습니다. 특히 커스텀 학습 루프가 필요한 경우나 복잡한 모델 구조를 다룰 때 매우 유용합니다.

핵심 포인트는 다음과 같습니다:

  • GradientTape를 사용한 자동미분 구현
  • 메모리 효율적인 대용량 데이터 처리
  • 그래디언트 클리핑과 정규화 기법 활용
  • 명확한 학습/평가 루프 분리