train_step() 재정의

가장 먼저 선택할 사용자 정의 지점

학습 방식을 바꿔야 할 때는 우선 train_step() 재정의를 검토해. 배치 하나의 순전파와 역전파를 담당하는 메서드만 교체하고 나머지는 그대로 물려받는 방식이야. 반복 전체는 여전히 fit()이 운전하므로 진행 표시줄, 콜백, 검증 단계, 분산 전략이 계속 작동해. Keras가 기본 학습 단계 대신 사용자가 작성한 단계를 호출할 뿐이지.

한 단계 안에서 처리할 네 가지

train_step()은 대체로 네 동작으로 이루어져. 배치를 풀고, 순전파로 y_pred를 만들고, self.compute_loss()로 손실을 구한 뒤, 그래디언트를 계산해 적용해. 마지막에는 평가지표 객체를 갱신하고 이름과 값을 담은 딕셔너리를 반환해. 바로 이 반환값이 fit()의 진행 표시줄에 나타나.

백엔드 차이가 드러나는 곳은 그래디언트 계산 방식이야. TensorFlow는 tf.GradientTape, PyTorch는 loss.backward() + optimizer.step(), JAX는 상태를 직접 넘기는 compute_loss_and_updates 형태를 사용해. Keras 3가 레이어 연산을 추상화해도 학습 단계의 그래디언트 처리는 각 백엔드 고유 방식에 맞춰야 해. 아래 코드는 TensorFlow 백엔드 버전이야.

Code

train_step() 재정의(TensorFlow 백엔드)·python

import tensorflow as tf
import keras

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

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

        # Compute and apply gradients (TF-native)
        grads = tape.gradient(loss, self.trainable_variables)
        self.optimizer.apply_gradients(
            zip(grads, self.trainable_variables)
        )

        # Update and return metrics
        for metric in self.metrics:
            if metric.name == "loss":
                metric.update_state(loss)
            else:
                metric.update_state(y, y_pred)

        return {m.name: m.result() for m in self.metrics}

# Still use fit()!
model = CustomModel(...)
model.compile(optimizer="adam", loss="mse")
model.fit(x_train, y_train, epochs=10)

가장 먼저 선택할 사용자 정의 지점

한 단계 안에서 처리할 네 가지

Code

External links

Exercise

Progress

댓글 0