TensorFlow에서의 `fit()` 동작을 커스터마이즈

import os

# 이 가이드는 TF 백엔드에서만 실행할 수 있습니다.
os.environ["KERAS_BACKEND"] = "tensorflow"

import tensorflow as tf
import keras
from keras import layers
import numpy as np

class CustomModel(keras.Model):
    def train_step(self, data):
        # 데이터를 언패킹합니다.
        # 데이터의 구조는 모델과 `fit()`에 전달하는 값에 따라 다릅니다.
        x, y = data

        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)  # 순전파
            # 손실 값을 계산합니다.
            # (손실 함수는 `compile()`에서 설정됩니다)
            loss = self.compute_loss(y=y, y_pred=y_pred)

        # 그래디언트를 계산합니다.
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)

        # 가중치를 업데이트합니다.
        self.optimizer.apply(gradients, trainable_vars)

        # 메트릭을 업데이트합니다. (손실을 추적하는 메트릭 포함)
        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}

# `CustomModel` 인스턴스를 생성하고 컴파일합니다.
inputs = keras.Input(shape=(32,))
outputs = keras.layers.Dense(1)(inputs)
model = CustomModel(inputs, outputs)
model.compile(optimizer="adam", loss="mse", metrics=["mae"])

# 평소처럼 `fit`을 사용하세요.
x = np.random.random((1000, 32))
y = np.random.random((1000, 1))
model.fit(x, y, epochs=3)

class CustomModel(keras.Model):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.loss_tracker = keras.metrics.Mean(name="loss")
        self.mae_metric = keras.metrics.MeanAbsoluteError(name="mae")
        self.loss_fn = keras.losses.MeanSquaredError()

    def train_step(self, data):
        x, y = data

        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)  # 순전파
            # 우리만의 손실을 계산합니다.
            loss = self.loss_fn(y, y_pred)

        # 그래디언트를 계산합니다.
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)

        # 가중치를 업데이트합니다.
        self.optimizer.apply(gradients, trainable_vars)

        # 우리만의 메트릭을 계산합니다.
        self.loss_tracker.update_state(loss)
        self.mae_metric.update_state(y, y_pred)
        return {
            "loss": self.loss_tracker.result(),
            "mae": self.mae_metric.result(),
        }

    @property
    def metrics(self):
        # 각 에포크의 시작 시점이나 `evaluate()`의 시작 시점에,
        # `reset_states()`가 자동으로 호출될 수 있도록,
        # `Metric` 객체를 여기에 나열합니다.
        return [self.loss_tracker, self.mae_metric]


# `CustomModel` 인스턴스를 생성하고 컴파일합니다.
inputs = keras.Input(shape=(32,))
outputs = keras.layers.Dense(1)(inputs)
model = CustomModel(inputs, outputs)

# 여기에서는 손실 함수나 메트릭을 전달하지 않습니다.
model.compile(optimizer="adam")

# 평소처럼 `fit`을 사용하세요 — 콜백 등을 사용할 수 있습니다.
x = np.random.random((1000, 32))
y = np.random.random((1000, 1))
model.fit(x, y, epochs=5)

class CustomModel(keras.Model):
    def train_step(self, data):
        # 데이터를 언팩합니다.
        # 데이터의 구조는 모델과 `fit()`에 전달하는 값에 따라 달라집니다.
        if len(data) == 3:
            x, y, sample_weight = data
        else:
            sample_weight = None
            x, y = data

        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)  # 순전파
            # 손실 값을 계산합니다.
            # (손실 함수는 `compile()`에서 설정됩니다)
            loss = self.compute_loss(
                y=y,
                y_pred=y_pred,
                sample_weight=sample_weight,
            )

        # 그래디언트를 계산합니다.
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)

        # 가중치를 업데이트합니다.
        self.optimizer.apply(gradients, trainable_vars)

        # 메트릭을 업데이트합니다.
        # 메트릭은 `compile()`에서 설정됩니다.
        for metric in self.metrics:
            if metric.name == "loss":
                metric.update_state(loss)
            else:
                metric.update_state(y, y_pred, sample_weight=sample_weight)

        # 메트릭 이름을 현재 값에 매핑하는 딕셔너리를 반환합니다.
        # 여기에는 손실(`self.metrics`에서 추적된)이 포함된다는 점에 유의하세요.
        return {m.name: m.result() for m in self.metrics}


# `CustomModel` 인스턴스를 생성하고 컴파일합니다.
inputs = keras.Input(shape=(32,))
outputs = keras.layers.Dense(1)(inputs)
model = CustomModel(inputs, outputs)
model.compile(optimizer="adam", loss="mse", metrics=["mae"])

# 이제 `sample_weight` 인자를 사용할 수 있습니다.
x = np.random.random((1000, 32))
y = np.random.random((1000, 1))
sw = np.random.random((1000, 1))
model.fit(x, y, sample_weight=sw, epochs=3)

class CustomModel(keras.Model):
    def test_step(self, data):
        # 데이터를 언팩합니다.
        x, y = data
        # 예측값을 계산합니다.
        y_pred = self(x, training=False)
        # 손실을 추적하는 메트릭을 업데이트합니다.
        loss = self.compute_loss(y=y, y_pred=y_pred)
        # 메트릭을 업데이트합니다.
        for metric in self.metrics:
            if metric.name == "loss":
                metric.update_state(loss)
            else:
                metric.update_state(y, y_pred)
        # 메트릭 이름을 현재 값에 매핑하는 딕셔너리를 반환합니다.
        # 여기에는 손실(`self.metrics`에서 추적된)이 포함된다는 점에 유의하세요.
        return {m.name: m.result() for m in self.metrics}


# `CustomModel` 인스턴스를 생성하고 컴파일합니다.
inputs = keras.Input(shape=(32,))
outputs = keras.layers.Dense(1)(inputs)
model = CustomModel(inputs, outputs)
model.compile(loss="mse", metrics=["mae"])

# 커스텀 `test_step`으로 평가합니다.
x = np.random.random((1000, 32))
y = np.random.random((1000, 1))
model.evaluate(x, y)

# 판별자를 생성합니다.
discriminator = keras.Sequential(
    [
        keras.Input(shape=(28, 28, 1)),
        layers.Conv2D(64, (3, 3), strides=(2, 2), padding="same"),
        layers.LeakyReLU(negative_slope=0.2),
        layers.Conv2D(128, (3, 3), strides=(2, 2), padding="same"),
        layers.LeakyReLU(negative_slope=0.2),
        layers.GlobalMaxPooling2D(),
        layers.Dense(1),
    ],
    name="discriminator",
)

# 생성자를 생성합니다.
latent_dim = 128
generator = keras.Sequential(
    [
        keras.Input(shape=(latent_dim,)),
        # 7x7x128 맵으로 reshape 할 128개의 계수를 생성하려고 합니다.
        layers.Dense(7 * 7 * 128),
        layers.LeakyReLU(negative_slope=0.2),
        layers.Reshape((7, 7, 128)),
        layers.Conv2DTranspose(128, (4, 4), strides=(2, 2), padding="same"),
        layers.LeakyReLU(negative_slope=0.2),
        layers.Conv2DTranspose(128, (4, 4), strides=(2, 2), padding="same"),
        layers.LeakyReLU(negative_slope=0.2),
        layers.Conv2D(1, (7, 7), padding="same", activation="sigmoid"),
    ],
    name="generator",
)

class GAN(keras.Model):
    def __init__(self, discriminator, generator, latent_dim):
        super().__init__()
        self.discriminator = discriminator
        self.generator = generator
        self.latent_dim = latent_dim
        self.d_loss_tracker = keras.metrics.Mean(name="d_loss")
        self.g_loss_tracker = keras.metrics.Mean(name="g_loss")
        self.seed_generator = keras.random.SeedGenerator(1337)

    @property
    def metrics(self):
        return [self.d_loss_tracker, self.g_loss_tracker]

    def compile(self, d_optimizer, g_optimizer, loss_fn):
        super().compile()
        self.d_optimizer = d_optimizer
        self.g_optimizer = g_optimizer
        self.loss_fn = loss_fn

    def train_step(self, real_images):
        if isinstance(real_images, tuple):
            real_images = real_images[0]
        # 잠재 공간에서 랜덤 포인트를 샘플링합니다.
        batch_size = tf.shape(real_images)[0]
        random_latent_vectors = keras.random.normal(
            shape=(batch_size, self.latent_dim), seed=self.seed_generator
        )

        # 그것들을 가짜 이미지로 디코딩합니다.
        generated_images = self.generator(random_latent_vectors)

        # 그것들을 진짜 이미지와 결합합니다.
        combined_images = tf.concat([generated_images, real_images], axis=0)

        # 진짜 이미지와 가짜 이미지를 구분하는 레이블을 구성합니다.
        labels = tf.concat(
            [tf.ones((batch_size, 1)), tf.zeros((batch_size, 1))], axis=0
        )
        # 레이블에 랜덤 노이즈를 추가합니다 - 중요한 트릭입니다!
        labels += 0.05 * keras.random.uniform(
            tf.shape(labels), seed=self.seed_generator
        )

        # 판별자를 트레이닝합니다.
        with tf.GradientTape() as tape:
            predictions = self.discriminator(combined_images)
            d_loss = self.loss_fn(labels, predictions)
        grads = tape.gradient(d_loss, self.discriminator.trainable_weights)
        self.d_optimizer.apply(grads, self.discriminator.trainable_weights)

        # 잠재 공간에서 랜덤 포인트를 샘플링합니다.
        random_latent_vectors = keras.random.normal(
            shape=(batch_size, self.latent_dim), seed=self.seed_generator
        )

        # '모두 진짜 이미지(all real images)'라는 레이블을 구성합니다.
        misleading_labels = tf.zeros((batch_size, 1))

        # 생성자를 트레이닝합니다!
        # (판별자의 가중치는 *업데이트하지 않아야* 한다는 점에 유의하세요)
        with tf.GradientTape() as tape:
            predictions = self.discriminator(self.generator(random_latent_vectors))
            g_loss = self.loss_fn(misleading_labels, predictions)
        grads = tape.gradient(g_loss, self.generator.trainable_weights)
        self.g_optimizer.apply(grads, self.generator.trainable_weights)

        # 메트릭을 업데이트하고 그 값을 반환합니다.
        self.d_loss_tracker.update_state(d_loss)
        self.g_loss_tracker.update_state(g_loss)
        return {
            "d_loss": self.d_loss_tracker.result(),
            "g_loss": self.g_loss_tracker.result(),
        }

# 데이터셋을 준비합니다. 트레이닝 및 테스트 모두 MNIST 숫자 데이터를 사용합니다.
batch_size = 64
(x_train, _), (x_test, _) = keras.datasets.mnist.load_data()
all_digits = np.concatenate([x_train, x_test])
all_digits = all_digits.astype("float32") / 255.0
all_digits = np.reshape(all_digits, (-1, 28, 28, 1))
dataset = tf.data.Dataset.from_tensor_slices(all_digits)
dataset = dataset.shuffle(buffer_size=1024).batch(batch_size)

gan = GAN(discriminator=discriminator, generator=generator, latent_dim=latent_dim)
gan.compile(
    d_optimizer=keras.optimizers.Adam(learning_rate=0.0003),
    g_optimizer=keras.optimizers.Adam(learning_rate=0.0003),
    loss_fn=keras.losses.BinaryCrossentropy(from_logits=True),
)

# 실행 시간을 제한하기 위해 100개의 배치에서만 트레이닝합니다.
# 전체 데이터셋으로 트레이닝할 수도 있습니다.
# 좋은 결과를 얻으려면 약 20 에포크가 필요합니다.
gan.fit(dataset.take(100), epochs=1)