Custom Training Loop 완벽 가이드

TensorFlow에서 모델 학습을 완전히 제어하는 Custom Training Loop의 모든 것을 다룹니다. 서브클래싱부터 분산 학습까지, 초급자도 이해할 수 있도록 실무 예제와 함께 설명합니다.

---

목차

1. tf.keras.Model 서브클래싱
2. train_step 커스터마이징
3. tf.function으로 성능 최적화
4. 그래디언트 클리핑
5. 혼합 정밀도 학습 (Mixed Precision)
6. 분산 학습 전략

---

1. tf.keras.Model 서브클래싱

김개발 씨는 TensorFlow로 딥러닝 모델을 만들던 중 고민에 빠졌습니다. Sequential API로는 원하는 복잡한 구조를 표현하기 어려웠기 때문입니다.

"더 자유롭게 모델을 설계할 수 있는 방법이 없을까요?"

tf.keras.Model 서브클래싱은 클래스를 상속받아 나만의 모델을 자유롭게 정의하는 방법입니다. 마치 레고 블록으로 원하는 형태를 직접 조립하는 것과 같습니다.

이 방식을 익히면 어떤 복잡한 아키텍처도 구현할 수 있게 됩니다.

다음 코드를 살펴봅시다.

import tensorflow as tf

class MyModel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        # 레이어 정의
        self.dense1 = tf.keras.layers.Dense(128, activation='relu')
        self.dropout = tf.keras.layers.Dropout(0.3)
        self.dense2 = tf.keras.layers.Dense(10)

    def call(self, inputs, training=False):
        # 순전파 로직 정의
        x = self.dense1(inputs)
        x = self.dropout(x, training=training)
        return self.dense2(x)

model = MyModel()

김개발 씨는 입사 후 첫 딥러닝 프로젝트를 맡게 되었습니다. 처음에는 Sequential API로 간단한 모델을 만들었습니다.

층을 쌓기만 하면 되니 정말 편리했습니다. 그런데 프로젝트 요구사항이 점점 복잡해졌습니다.

중간에 분기가 필요하고, 여러 입력을 받아야 하고, 특정 조건에 따라 다른 연산을 수행해야 했습니다. Sequential API로는 도저히 표현할 수 없는 구조였습니다.

선배 개발자 박시니어 씨가 조언했습니다. "그럴 땐 서브클래싱을 사용하면 돼요.

파이썬 클래스처럼 모델을 정의할 수 있거든요." 서브클래싱이란 무엇일까요? 쉽게 비유하자면, 기성품 가구와 맞춤 제작 가구의 차이와 같습니다.

Sequential API는 이미 만들어진 가구를 순서대로 배치하는 것입니다. 반면 서브클래싱은 목공 도구를 직접 들고 원하는 가구를 처음부터 만드는 것입니다.

tf.keras.Model 클래스를 상속받으면 두 가지 핵심 메서드를 정의해야 합니다. 첫 번째는 init 메서드입니다.

여기서 모델에 사용할 모든 레이어를 정의합니다. Dense, Conv2D, LSTM 등 필요한 레이어를 인스턴스 변수로 생성해둡니다.

두 번째는 call 메서드입니다. 이 메서드가 실제 순전파 로직을 담당합니다.

입력 데이터가 어떤 레이어를 거쳐 어떻게 변환되는지를 여기서 정의합니다. if문이나 for문 같은 파이썬 제어문도 자유롭게 사용할 수 있습니다.

위 코드를 살펴보겠습니다. __init__에서 Dense 레이어 두 개와 Dropout 레이어를 정의했습니다.

call 메서드에서는 입력이 dense1을 거치고, dropout을 통과한 뒤, 마지막으로 dense2를 거쳐 출력됩니다. training 매개변수에 주목하세요.

Dropout이나 BatchNormalization 같은 레이어는 학습할 때와 추론할 때 다르게 동작합니다. training 매개변수로 이를 제어할 수 있습니다.

실무에서 서브클래싱은 정말 자주 사용됩니다. ResNet, Transformer 같은 복잡한 아키텍처를 구현할 때 필수입니다.

논문에 나온 새로운 모델 구조를 직접 구현해야 할 때도 서브클래싱이 답입니다. 주의할 점도 있습니다.

서브클래싱 모델은 build 메서드가 호출되기 전까지 가중치가 생성되지 않습니다. 따라서 model.summary()를 바로 호출하면 에러가 발생할 수 있습니다.

먼저 더미 데이터로 한 번 호출해주면 해결됩니다. 김개발 씨는 서브클래싱을 익힌 후 복잡한 모델도 자신 있게 구현하게 되었습니다.

코드가 더 명확해지고, 디버깅도 쉬워졌습니다.

실전 팁

💡 - super().init() 호출을 잊지 마세요. 부모 클래스 초기화가 필수입니다.

• call 메서드에서 training 매개변수를 활용하면 학습/추론 모드를 구분할 수 있습니다.

---

2. train step 커스터마이징

김개발 씨가 model.fit()으로 학습을 돌리던 중 의문이 생겼습니다. "학습 과정에서 특별한 로직을 추가하고 싶은데, fit() 안에서는 어떻게 수정하지?" 박시니어 씨가 웃으며 말했습니다.

"train_step을 오버라이드하면 돼요."

train_step은 model.fit() 호출 시 매 배치마다 실행되는 학습 로직입니다. 이 메서드를 오버라이드하면 손실 계산, 그래디언트 업데이트 등을 완전히 제어할 수 있습니다.

마치 자동차의 엔진을 직접 튜닝하는 것과 같습니다.

다음 코드를 살펴봅시다.

class CustomModel(tf.keras.Model):
    def train_step(self, data):
        x, y = data

        with tf.GradientTape() as tape:
            # 순전파
            y_pred = self(x, training=True)
            # 손실 계산
            loss = self.compute_loss(y=y, y_pred=y_pred)

        # 그래디언트 계산 및 적용
        gradients = tape.gradient(loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))

        # 메트릭 업데이트
        for metric in self.metrics:
            metric.update_state(y, y_pred)
        return {m.name: m.result() for m in self.metrics}

김개발 씨는 GAN 모델을 학습시켜야 했습니다. 생성자와 판별자를 번갈아 학습시키는 특수한 로직이 필요했습니다.

일반적인 model.fit()으로는 이런 복잡한 학습 과정을 구현할 수 없었습니다. 박시니어 씨가 설명을 시작했습니다.

"model.fit()을 호출하면 내부적으로 train_step이 반복 실행돼요. 이 메서드만 오버라이드하면 fit()의 편리함은 유지하면서 학습 로직은 마음대로 바꿀 수 있어요." train_step의 동작을 이해하려면 먼저 GradientTape를 알아야 합니다.

GradientTape는 마치 녹화기와 같습니다. tape.gradient()를 호출하면 녹화된 연산들을 역으로 추적하며 그래디언트를 계산합니다.

with tf.GradientTape() as tape 블록 안에서 일어나는 모든 연산이 기록됩니다. 순전파를 수행하고 손실을 계산하는 과정이 모두 녹화됩니다.

블록을 벗어나면 tape.gradient()로 기록된 연산의 그래디언트를 구할 수 있습니다. 코드를 단계별로 살펴보겠습니다.

먼저 data에서 입력 x와 레이블 y를 분리합니다. GradientTape 컨텍스트 안에서 self(x, training=True)로 예측값을 얻습니다.

여기서 self를 호출하면 앞서 정의한 call 메서드가 실행됩니다. 다음으로 compute_loss로 손실을 계산합니다.

이 메서드는 compile() 시 지정한 손실 함수를 사용합니다. 물론 직접 손실을 계산해도 됩니다.

그래디언트 계산이 핵심입니다. tape.gradient(loss, self.trainable_variables)는 손실에 대한 모든 학습 가능한 변수의 그래디언트를 반환합니다.

마치 각 가중치가 손실에 얼마나 기여했는지 계산하는 것입니다. optimizer.apply_gradients()가 실제로 가중치를 업데이트합니다.

그래디언트와 변수를 zip으로 묶어 전달하면 옵티마이저가 학습률에 맞게 가중치를 조정합니다. 마지막으로 메트릭을 업데이트하고 결과를 반환합니다.

이 반환값이 fit() 실행 중 화면에 출력되는 값들입니다. 실무에서 train_step 커스터마이징은 다양하게 활용됩니다.

GAN처럼 여러 모델을 번갈아 학습할 때, 특수한 정규화를 적용할 때, 멀티태스크 학습에서 손실을 다르게 가중할 때 등 활용 범위가 넓습니다. 김개발 씨는 train_step을 오버라이드하여 GAN 학습 로직을 성공적으로 구현했습니다.

fit()의 콜백, 프로그레스 바 같은 편의 기능도 그대로 사용할 수 있어서 매우 만족스러웠습니다.

실전 팁

💡 - test_step도 오버라이드하면 evaluate() 동작도 커스터마이징할 수 있습니다.

• 여러 손실을 합칠 때는 self.add_loss()를 활용하면 깔끔합니다.

---

3. tf.function으로 성능 최적화

학습 코드를 완성한 김개발 씨가 실행 버튼을 눌렀습니다. 그런데 예상보다 학습이 느렸습니다.

"GPU를 쓰는데 왜 이렇게 느리지?" 박시니어 씨가 코드를 보더니 말했습니다. "tf.function 데코레이터를 붙여봐요."

tf.function은 파이썬 함수를 TensorFlow 그래프로 컴파일하는 데코레이터입니다. 마치 통역사 없이 직접 대화하는 것처럼, 파이썬 인터프리터를 거치지 않고 최적화된 연산을 수행합니다.

학습 속도가 눈에 띄게 향상됩니다.

다음 코드를 살펴봅시다.

class OptimizedModel(tf.keras.Model):
    @tf.function
    def train_step(self, data):
        x, y = data

        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)
            loss = self.compute_loss(y=y, y_pred=y_pred)

        # 그래프 모드에서 최적화된 연산 수행
        gradients = tape.gradient(loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))

        for metric in self.metrics:
            metric.update_state(y, y_pred)
        return {m.name: m.result() for m in self.metrics}

TensorFlow에는 두 가지 실행 모드가 있습니다. Eager 모드와 Graph 모드입니다.

TensorFlow 2.0부터 기본값은 Eager 모드입니다. 코드를 작성하는 즉시 실행되어 디버깅이 편리합니다.

하지만 Eager 모드에는 단점이 있습니다. 매 연산마다 파이썬 인터프리터를 거쳐야 합니다.

마치 외국어 대화를 할 때 매 문장마다 통역사를 거치는 것과 같습니다. 오버헤드가 쌓이면 성능 저하로 이어집니다.

Graph 모드는 다릅니다. 연산 전체를 하나의 그래프로 컴파일합니다.

통역사 없이 직접 대화하는 것처럼 빠릅니다. 게다가 TensorFlow가 그래프를 분석해서 불필요한 연산을 제거하고, 병렬 처리가 가능한 연산을 자동으로 병렬화합니다.

@tf.function 데코레이터를 붙이면 해당 함수가 처음 호출될 때 그래프로 컴파일됩니다. 이 과정을 트레이싱이라고 합니다.

한 번 트레이싱되면 이후 호출은 컴파일된 그래프로 실행됩니다. 코드는 이전과 거의 동일합니다.

단지 @tf.function 한 줄만 추가했을 뿐입니다. 하지만 내부적으로는 완전히 다르게 동작합니다.

파이썬 코드가 TensorFlow 연산 그래프로 변환되어 실행됩니다. 실제로 얼마나 빨라질까요?

모델과 데이터에 따라 다르지만, 일반적으로 2배에서 5배 정도 속도 향상을 기대할 수 있습니다. 특히 GPU를 사용할 때 효과가 극대화됩니다.

GPU와 CPU 사이의 데이터 전송이 최적화되기 때문입니다. 주의할 점이 있습니다.

tf.function 내부에서는 파이썬 부작용이 예상대로 동작하지 않을 수 있습니다. print() 문은 트레이싱 시에만 실행됩니다.

리스트에 값을 추가하는 것 같은 파이썬 객체 조작도 문제가 될 수 있습니다. 디버깅할 때는 tf.function을 잠시 제거하는 것이 좋습니다.

Eager 모드에서 먼저 코드가 정상 동작하는지 확인한 후, 성능 최적화를 위해 tf.function을 적용하세요. 또 하나 알아둘 점이 있습니다.

입력 데이터의 shape이나 dtype이 바뀌면 함수가 다시 트레이싱됩니다. 이를 리트레이싱이라고 합니다.

잦은 리트레이싱은 오히려 성능을 저하시킵니다. input_signature를 지정하면 리트레이싱을 방지할 수 있습니다.

김개발 씨는 @tf.function을 적용한 후 학습 시간이 절반으로 줄어든 것을 확인했습니다. 같은 코드인데 한 줄 추가로 이런 차이가 나다니, 정말 놀라웠습니다.

실전 팁

💡 - 디버깅 시에는 tf.function을 제거하고 Eager 모드로 테스트하세요.

• tf.print()를 사용하면 그래프 모드에서도 출력이 가능합니다.

---

4. 그래디언트 클리핑

김개발 씨가 RNN 모델을 학습시키던 중 이상한 현상을 발견했습니다. 손실 값이 갑자기 NaN이 되어버린 것입니다.

"왜 갑자기 학습이 망가지는 거죠?" 박시니어 씨가 진단했습니다. "그래디언트 폭발이에요.

클리핑을 적용해야 해요."

그래디언트 클리핑은 그래디언트의 크기를 제한하는 기법입니다. 마치 과속 방지턱처럼, 그래디언트가 너무 커지는 것을 막아줍니다.

RNN이나 Transformer 같은 깊은 모델에서 학습 안정성을 크게 높여줍니다.

다음 코드를 살펴봅시다.

class ClippedModel(tf.keras.Model):
    def __init__(self, clip_norm=1.0):
        super().__init__()
        self.clip_norm = clip_norm
        self.lstm = tf.keras.layers.LSTM(128, return_sequences=True)
        self.dense = tf.keras.layers.Dense(10)

    def train_step(self, data):
        x, y = data
        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)
            loss = self.compute_loss(y=y, y_pred=y_pred)

        gradients = tape.gradient(loss, self.trainable_variables)
        # 그래디언트 클리핑 적용
        gradients, _ = tf.clip_by_global_norm(gradients, self.clip_norm)
        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
        return {"loss": loss}

딥러닝에서 그래디언트는 손실 함수의 기울기입니다. 이 기울기 방향으로 가중치를 조금씩 업데이트하며 학습이 진행됩니다.

그런데 때때로 그래디언트가 비정상적으로 커지는 현상이 발생합니다. 이를 그래디언트 폭발이라고 합니다.

왜 그래디언트 폭발이 일어날까요? 역전파 과정을 생각해보세요.

깊은 네트워크에서는 그래디언트가 여러 층을 거슬러 올라갑니다. 각 층에서 그래디언트가 조금씩 증폭될 수 있습니다.

100개 층을 거치면서 매번 1.1배씩 증폭된다면? 결과는 기하급수적으로 커집니다.

특히 RNN 계열 모델이 취약합니다. 시퀀스 길이가 길어질수록 같은 가중치 행렬이 반복해서 곱해집니다.

행렬의 고유값이 1보다 크면 그래디언트가 폭발합니다. 그 결과 가중치가 NaN이 되고 학습이 완전히 망가집니다.

그래디언트 클리핑은 이 문제의 해결책입니다. 비유하자면 과속 방지턱과 같습니다.

자동차가 아무리 빨라도 방지턱을 넘으면 속도가 제한됩니다. 마찬가지로 그래디언트가 아무리 커도 클리핑 임계값을 넘으면 잘려나갑니다.

클리핑 방법에는 여러 가지가 있습니다. clip_by_value는 각 그래디언트 요소를 개별적으로 잘라냅니다.

clip_by_norm은 각 그래디언트 텐서의 노름을 제한합니다. clip_by_global_norm은 모든 그래디언트를 하나의 벡터로 보고 전체 노름을 제한합니다.

가장 널리 사용되는 것은 clip_by_global_norm입니다. 그래디언트들의 상대적 비율을 유지하면서 전체 크기만 제한하기 때문입니다.

논문에서 말하는 그래디언트 클리핑은 대부분 이 방법을 의미합니다. 코드를 보면 gradients = tape.gradient() 이후에 클리핑을 적용합니다.

tf.clip_by_global_norm은 두 값을 반환합니다. 첫 번째는 클리핑된 그래디언트 리스트, 두 번째는 원래 전체 노름입니다.

보통 첫 번째 값만 사용합니다. clip_norm 값은 어떻게 정할까요?

일반적으로 1.0에서 5.0 사이를 사용합니다. 학습 초기에 그래디언트 노름을 모니터링하고, 적절한 값을 선택하는 것이 좋습니다.

너무 작으면 학습이 느려지고, 너무 크면 클리핑 효과가 없습니다. Transformer, BERT, GPT 같은 대형 언어 모델에서도 그래디언트 클리핑은 필수입니다.

학습 안정성을 위해 거의 항상 적용됩니다. 클리핑 없이 학습하면 중간에 발산하는 경우가 많습니다.

김개발 씨는 clip_norm=1.0을 적용한 후 NaN 문제가 사라졌습니다. 학습이 안정적으로 수렴하는 것을 보며 안도의 한숨을 내쉬었습니다.

실전 팁

💡 - RNN, LSTM, Transformer 학습 시 클리핑은 거의 필수입니다.

• 그래디언트 노름을 텐서보드로 모니터링하면 적절한 clip_norm 값을 찾기 쉽습니다.

---

5. 혼합 정밀도 학습 (Mixed Precision)

김개발 씨가 대용량 모델을 학습시키려는데 GPU 메모리가 부족했습니다. 배치 크기를 줄이니 학습이 너무 느렸습니다.

"메모리도 아끼고 속도도 높이는 방법 없을까요?" 박시니어 씨가 답했습니다. "혼합 정밀도 학습을 써보세요."

혼합 정밀도 학습은 32비트와 16비트 부동소수점을 섞어 사용하는 기법입니다. 마치 고속도로에서 빠른 차선과 안전한 차선을 적절히 활용하는 것과 같습니다.

메모리 사용량은 줄이고 연산 속도는 높이면서도 정확도는 유지합니다.

다음 코드를 살펴봅시다.

# 혼합 정밀도 정책 설정
policy = tf.keras.mixed_precision.Policy('mixed_float16')
tf.keras.mixed_precision.set_global_policy(policy)

class MixedPrecisionModel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.dense1 = tf.keras.layers.Dense(256, activation='relu')
        self.dense2 = tf.keras.layers.Dense(10)

    def call(self, inputs):
        x = self.dense1(inputs)
        x = self.dense2(x)
        # 출력을 float32로 캐스팅 (손실 계산 안정성)
        return tf.cast(x, tf.float32)

# LossScaleOptimizer로 언더플로우 방지
optimizer = tf.keras.optimizers.Adam(1e-4)

컴퓨터에서 숫자를 표현하는 방식에는 여러 가지가 있습니다. float32는 32비트를 사용해 소수점 숫자를 표현합니다.

정밀도가 높지만 메모리를 많이 차지합니다. float16은 16비트만 사용합니다.

메모리는 절반이지만 표현 가능한 범위가 좁습니다. 딥러닝에서는 전통적으로 float32를 사용해왔습니다.

그래디언트 계산의 정밀도를 보장하기 위해서입니다. 하지만 모델이 점점 커지면서 메모리 문제가 심각해졌습니다.

GPT-3 같은 초거대 모델은 float32로는 도저히 학습할 수 없습니다. 혼합 정밀도는 두 세계의 장점을 결합합니다.

순전파와 역전파의 대부분 연산은 float16으로 수행합니다. 빠르고 메모리 효율적입니다.

하지만 가중치 저장과 그래디언트 누적은 float32로 합니다. 정밀도가 중요한 부분만 안전하게 처리하는 것입니다.

비유하자면 이렇습니다. 고속도로에서 빠른 차선은 speed를 위해, 느린 차선은 safety를 위해 사용합니다.

혼합 정밀도도 마찬가지입니다. 연산 속도가 중요한 곳은 float16으로, 정밀도가 중요한 곳은 float32로 처리합니다.

TensorFlow에서 혼합 정밀도를 사용하는 방법은 간단합니다. mixed_float16 정책을 설정하면 됩니다.

그러면 레이어들이 자동으로 float16 연산을 수행합니다. 개발자가 직접 캐스팅할 필요가 거의 없습니다.

한 가지 주의할 점이 있습니다. float16은 표현 범위가 좁아서 언더플로우가 발생할 수 있습니다.

아주 작은 그래디언트가 0이 되어버리는 것입니다. 이를 방지하기 위해 Loss Scaling을 사용합니다.

손실에 큰 수를 곱해서 그래디언트를 키우고, 업데이트 전에 다시 나눕니다. TensorFlow 2.x에서는 optimizer가 자동으로 loss scaling을 처리합니다.

mixed_float16 정책을 사용하면 옵티마이저가 알아서 스케일링을 적용합니다. 개발자가 신경 쓸 것이 거의 없습니다.

성능 향상은 얼마나 될까요? NVIDIA의 Tensor Core가 있는 GPU에서 2배에서 3배 속도 향상을 기대할 수 있습니다.

메모리 사용량도 거의 절반으로 줄어듭니다. 같은 GPU로 더 큰 배치 크기, 더 큰 모델을 학습할 수 있습니다.

출력층에서 float32로 캐스팅하는 것을 잊지 마세요. 손실 계산은 float32로 해야 수치적으로 안정적입니다.

모델의 마지막에 tf.cast(x, tf.float32)를 추가하는 것이 좋습니다. 김개발 씨는 혼합 정밀도를 적용한 후 배치 크기를 두 배로 늘릴 수 있었습니다.

학습 시간도 40% 단축되었습니다. GPU 활용도가 눈에 띄게 올라간 것을 확인하고 뿌듯해했습니다.

실전 팁

💡 - NVIDIA GPU의 Tensor Core를 활용하려면 배치 크기와 레이어 크기를 8의 배수로 맞추세요.

• 손실이 NaN이 되면 loss scale 설정을 조정해보세요.

---

6. 분산 학습 전략

회사에서 GPU 서버를 여러 대 지원받게 된 김개발 씨. 하지만 여러 GPU를 어떻게 활용해야 할지 막막했습니다.

"GPU 4개를 동시에 쓰려면 코드를 완전히 다시 짜야 하나요?" 박시니어 씨가 웃으며 말했습니다. "tf.distribute.Strategy를 쓰면 몇 줄만 추가하면 돼요."

분산 학습 전략은 여러 GPU나 여러 머신에서 모델을 병렬로 학습시키는 방법입니다. 마치 여러 명의 요리사가 동시에 요리하는 것처럼, 데이터를 나눠 처리하고 결과를 합칩니다.

학습 시간을 GPU 개수만큼 단축할 수 있습니다.

다음 코드를 살펴봅시다.

# 멀티 GPU 분산 전략 설정
strategy = tf.distribute.MirroredStrategy()
print(f'Number of devices: {strategy.num_replicas_in_sync}')

with strategy.scope():
    # 전략 스코프 안에서 모델과 옵티마이저 생성
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(128, activation='relu'),
        tf.keras.layers.Dense(10)
    ])

    optimizer = tf.keras.optimizers.Adam(0.001)
    model.compile(optimizer=optimizer,
                  loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                  metrics=['accuracy'])

# 배치 크기는 GPU 개수에 비례하여 증가
BATCH_SIZE = 64 * strategy.num_replicas_in_sync
model.fit(train_dataset.batch(BATCH_SIZE), epochs=10)

딥러닝 모델이 커지면서 학습 시간도 기하급수적으로 늘어났습니다. ImageNet 데이터셋으로 ResNet을 학습하는 데 단일 GPU로는 일주일 넘게 걸립니다.

하지만 8개의 GPU를 사용하면 하루 만에 끝낼 수 있습니다. 이것이 분산 학습의 힘입니다.

분산 학습에는 크게 두 가지 방식이 있습니다. 데이터 병렬화와 모델 병렬화입니다.

데이터 병렬화는 같은 모델을 여러 GPU에 복제하고, 데이터를 나눠 처리합니다. 모델 병렬화는 하나의 모델을 여러 GPU에 나눠 배치합니다.

가장 일반적으로 사용되는 것은 데이터 병렬화입니다. TensorFlow의 tf.distribute.Strategy API는 분산 학습을 놀라울 정도로 쉽게 만들어줍니다.

기존 코드를 거의 수정하지 않고도 분산 학습을 적용할 수 있습니다. 마법처럼 느껴질 정도입니다.

MirroredStrategy가 가장 기본적인 전략입니다. 한 머신 안의 여러 GPU를 활용합니다.

각 GPU에 모델의 복사본을 만들고, 배치 데이터를 GPU 개수만큼 나눕니다. 각 GPU가 자신의 데이터로 그래디언트를 계산하고, 모든 그래디언트를 합쳐서 가중치를 업데이트합니다.

이 과정을 All-Reduce라고 합니다. 마치 여러 명의 학생이 각자 문제를 풀고, 답을 모아서 평균을 내는 것과 같습니다.

NVIDIA의 NCCL 라이브러리가 GPU 간 통신을 최적화합니다. 코드를 살펴보면, strategy.scope() 컨텍스트 안에서 모델과 옵티마이저를 생성합니다.

이 한 줄이 핵심입니다. 스코프 안에서 생성된 모든 변수가 자동으로 각 GPU에 복제됩니다.

그래디언트 동기화도 자동으로 처리됩니다. 배치 크기 설정에 주의가 필요합니다.

분산 학습에서 글로벌 배치 크기는 GPU 개수에 비례해서 증가시키는 것이 일반적입니다. GPU 4개를 사용한다면 배치 크기를 4배로 늘립니다.

각 GPU가 원래 배치 크기만큼 처리하기 때문입니다. 다른 전략들도 있습니다.

MultiWorkerMirroredStrategy는 여러 머신에 걸쳐 분산 학습을 수행합니다. TPUStrategy는 Google의 TPU를 활용합니다.

ParameterServerStrategy는 파라미터 서버 아키텍처를 사용합니다. 상황에 맞는 전략을 선택하면 됩니다.

실무에서 분산 학습을 적용할 때 몇 가지 팁이 있습니다. 먼저 학습률을 조정해야 할 수 있습니다.

배치 크기가 커지면 학습률도 높여야 할 때가 많습니다. Linear Scaling Rule에 따르면 배치 크기가 n배가 되면 학습률도 n배로 늘립니다.

또한 학습률 워밍업을 적용하는 것이 좋습니다. 처음에는 낮은 학습률로 시작해서 점진적으로 높입니다.

큰 배치로 학습할 때 초기 불안정성을 줄여줍니다. 김개발 씨는 4개의 GPU에 MirroredStrategy를 적용했습니다.

학습 시간이 거의 4분의 1로 줄었습니다. 코드 변경은 불과 몇 줄이었습니다.

"이렇게 쉬울 줄이야!" 김개발 씨는 감탄했습니다.

실전 팁

💡 - 글로벌 배치 크기가 커지면 학습률도 비례해서 조정하세요.

• 학습률 워밍업을 적용하면 큰 배치 학습의 안정성이 높아집니다.

---

이상으로 학습을 마칩니다. 위 내용을 직접 코드로 작성해보면서 익혀보세요!