사용자 정의 학습 단계

fit()의 기반 기능은 두고 한 단계만 바꿔

model.fit()을 그대로 쓰는 것과 학습 루프를 전부 직접 작성하는 것 사이에는 train_step() 재정의라는 선택지가 있어. 배치마다 수행할 손실과 기울기 계산만 바꾸면서 진행 표시줄, EarlyStopping·ModelCheckpoint 콜백, 검증, 분산 학습은 계속 fit()에 맡길 수 있지. GAN, 지식 증류, 커리큘럼 학습처럼 일반적인 손실과 역전파로 부족할 때 알맞아.

train_step()이 지켜야 할 두 가지

train_step(self, data)는 배치 하나를 받아 실제로 가중치를 갱신해야 해. 손실을 계산하고 self.trainable_variables에 대한 기울기를 구한 뒤 self.optimizer로 적용해. 또한 평가지표 이름과 값을 담은 딕셔너리를 반환해야 fit()이 진행 표시줄과 콜백에 전달할 수 있어. 손실 함수를 직접 호출하기보다 self.compute_loss()를 쓰면 정규화 손실과 compile(loss=...) 설정도 함께 적용돼. 가중치 갱신을 빼먹으면 모델은 학습하지 않고, 반환 딕셔너리가 틀리면 콜백이 상태를 볼 수 없어.

Code

fit()을 유지하며 train_step() 재정의·python

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

        # Forward pass with gradient tracking
        y_pred = self(x, training=True)
        loss = self.compute_loss(y=y, y_pred=y_pred)

        # Compute and apply gradients
        gradients = self.optimizer.compute_gradients(loss, self.trainable_variables)
        self.optimizer.apply(gradients)

        # Update 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}

fit()의 기반 기능은 두고 한 단계만 바꿔

train_step()이 지켜야 할 두 가지

Code

External links

Exercise

Progress

댓글 0