Custom training step

fit() 은 그대로, step 만 바꿔

'model.fit() 그대로' 와 '학습 loop 를 처음부터 손으로' 사이에 한 단계의 커스터마이즈가 있어. train_step() 을 override 하면 batch 당 로직 — gradient 계산 — 만 갈아끼우고, fit() 이 공짜로 주는 건 다 유지해: progress bar, EarlyStopping / ModelCheckpoint 같은 callback, validation pass, distribution strategy. GAN, knowledge distillation, curriculum learning, 또는 평범한 loss-and-backprop step 으로 부족한 모든 경우에 이게 맞는 도구야 (Code block 참고).

train_step() 이 지켜야 할 계약

train_step(self, data) 는 batch 하나를 받고, Keras 는 두 가지를 기대해. 첫째, 메서드 안에서 실제로 weight 를 업데이트해야 해 — loss 계산, self.trainable_variables 에 대한 gradient 도출, self.optimizer 로 적용. 둘째, metric 이름 → 값 dict 를 반환해야 해. 그 dict 가 바로 fit() 이 progress bar 에 찍고 모든 callback 에 먹이는 거야. loss 함수를 직접 부르지 말고 self.compute_loss() 를 써 — regularization 을 같이 접어주고 compile(loss=...) 가 계속 돌게 해줘. dict 를 틀리면 callback 이 아무것도 못 보고, optimizer apply 를 빼먹으면 model 이 아무것도 학습 안 해.

Code

fit() 유지하면서 train_step() override·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() 은 그대로, step 만 바꿔

train_step() 이 지켜야 할 계약

Code

External links

Exercise

Progress

댓글 0