TensorFlow에서 처음부터 트레이닝 루프 작성하기

import time
import os

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

import tensorflow as tf
import keras
import numpy as np

def get_model():
    inputs = keras.Input(shape=(784,), name="digits")
    x1 = keras.layers.Dense(64, activation="relu")(inputs)
    x2 = keras.layers.Dense(64, activation="relu")(x1)
    outputs = keras.layers.Dense(10, name="predictions")(x2)
    model = keras.Model(inputs=inputs, outputs=outputs)
    return model


model = get_model()

# 옵티마이저 인스턴스화
optimizer = keras.optimizers.Adam(learning_rate=1e-3)
# 손실 함수 인스턴스화
loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# 트레이닝 데이터셋 준비
batch_size = 32
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = np.reshape(x_train, (-1, 784))
x_test = np.reshape(x_test, (-1, 784))

# 검증을 위해 10,000개의 샘플을 예약합니다.
x_val = x_train[-10000:]
y_val = y_train[-10000:]
x_train = x_train[:-10000]
y_train = y_train[:-10000]

# 트레이닝 데이터셋 준비
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = train_dataset.shuffle(buffer_size=1024).batch(batch_size)

# 검증 데이터셋 준비
val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))
val_dataset = val_dataset.batch(batch_size)

epochs = 3
for epoch in range(epochs):
    print(f"\nStart of epoch {epoch}")

    # 데이터셋의 배치에 걸쳐 반복합니다.
    for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):
        # 순전파 동안 실행된 연산을 기록하기 위해, GradientTape를 엽니다.
        # 이를 통해 자동 미분이 가능합니다.
        with tf.GradientTape() as tape:
            # 레이어의 순전파를 실행합니다.
            # 레이어가 입력에 적용하는 연산은 GradientTape에 기록됩니다.
            logits = model(x_batch_train, training=True)  # 이 미니배치에 대한 로짓

            # 이 미니배치에 대한 손실 값을 계산합니다.
            loss_value = loss_fn(y_batch_train, logits)

        # 그레디언트 테이프를 사용해 손실에 대한
        # 트레이닝 가능한 변수의 그래디언트를 자동으로 가져옵니다.
        grads = tape.gradient(loss_value, model.trainable_weights)

        # 그래디언트 하강법의 한 단계를 실행하여
        # 손실을 최소화하기 위해 변수의 값을 업데이트합니다.
        optimizer.apply(grads, model.trainable_weights)

        # 100 배치마다 로그를 출력합니다.
        if step % 100 == 0:
            print(
                f"Training loss (for 1 batch) at step {step}: {float(loss_value):.4f}"
            )
            print(f"Seen so far: {(step + 1) * batch_size} samples")

# 새로운 모델 가져오기
model = get_model()

# 모델을 트레이닝할 옵티마이저 인스턴스화
optimizer = keras.optimizers.Adam(learning_rate=1e-3)
# 손실 함수 인스턴스화
loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# 메트릭 준비
train_acc_metric = keras.metrics.SparseCategoricalAccuracy()
val_acc_metric = keras.metrics.SparseCategoricalAccuracy()

epochs = 2
for epoch in range(epochs):
    print(f"\nStart of epoch {epoch}")
    start_time = time.time()

    # 데이터셋의 배치에 걸쳐 반복합니다.
    for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):
        with tf.GradientTape() as tape:
            logits = model(x_batch_train, training=True)
            loss_value = loss_fn(y_batch_train, logits)
        grads = tape.gradient(loss_value, model.trainable_weights)
        optimizer.apply(grads, model.trainable_weights)

        # 트레이닝 메트릭 업데이트.
        train_acc_metric.update_state(y_batch_train, logits)

        # 100 배치마다 로그를 출력합니다.
        if step % 100 == 0:
            print(
                f"Training loss (for 1 batch) at step {step}: {float(loss_value):.4f}"
            )
            print(f"Seen so far: {(step + 1) * batch_size} samples")

    # 각 에포크가 끝날 때 메트릭을 표시합니다.
    train_acc = train_acc_metric.result()
    print(f"Training acc over epoch: {float(train_acc):.4f}")

    # 각 에포크가 끝날 때 트레이닝 메트릭을 초기화합니다.
    train_acc_metric.reset_state()

    # 각 에포크가 끝날 때 검증 루프를 실행합니다.
    for x_batch_val, y_batch_val in val_dataset:
        val_logits = model(x_batch_val, training=False)
        # 검증 메트릭 업데이트
        val_acc_metric.update_state(y_batch_val, val_logits)
    val_acc = val_acc_metric.result()
    val_acc_metric.reset_state()
    print(f"Validation acc: {float(val_acc):.4f}")
    print(f"Time taken: {time.time() - start_time:.2f}s")

@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        logits = model(x, training=True)
        loss_value = loss_fn(y, logits)
    grads = tape.gradient(loss_value, model.trainable_weights)
    optimizer.apply(grads, model.trainable_weights)
    train_acc_metric.update_state(y, logits)
    return loss_value

@tf.function
def test_step(x, y):
    val_logits = model(x, training=False)
    val_acc_metric.update_state(y, val_logits)

epochs = 2
for epoch in range(epochs):
    print(f"\nStart of epoch {epoch}")
    start_time = time.time()

    # 데이터셋의 배치에 걸쳐 반복합니다.
    for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):
        loss_value = train_step(x_batch_train, y_batch_train)

        # 100 배치마다 로그를 출력합니다.
        if step % 100 == 0:
            print(
                f"Training loss (for 1 batch) at step {step}: {float(loss_value):.4f}"
            )
            print(f"Seen so far: {(step + 1) * batch_size} samples")

    # 각 에포크가 끝날 때 메트릭을 표시합니다.
    train_acc = train_acc_metric.result()
    print(f"Training acc over epoch: {float(train_acc):.4f}")

    # 각 에포크가 끝날 때 트레이닝 메트릭을 초기화합니다.
    train_acc_metric.reset_state()

    # 각 에포크가 끝날 때 검증 루프를 실행합니다.
    for x_batch_val, y_batch_val in val_dataset:
        test_step(x_batch_val, y_batch_val)

    val_acc = val_acc_metric.result()
    val_acc_metric.reset_state()
    print(f"Validation acc: {float(val_acc):.4f}")
    print(f"Time taken: {time.time() - start_time:.2f}s")

class ActivityRegularizationLayer(keras.layers.Layer):
    def call(self, inputs):
        self.add_loss(1e-2 * tf.reduce_sum(inputs))
        return inputs

inputs = keras.Input(shape=(784,), name="digits")
x = keras.layers.Dense(64, activation="relu")(inputs)
# activity 정규화 레이어 삽입
x = ActivityRegularizationLayer()(x)
x = keras.layers.Dense(64, activation="relu")(x)
outputs = keras.layers.Dense(10, name="predictions")(x)

model = keras.Model(inputs=inputs, outputs=outputs)

@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        logits = model(x, training=True)
        loss_value = loss_fn(y, logits)
        # 순전파 동안 생성된 추가 손실을 더합니다.
        loss_value += sum(model.losses)
    grads = tape.gradient(loss_value, model.trainable_weights)
    optimizer.apply(grads, model.trainable_weights)
    train_acc_metric.update_state(y, logits)
    return loss_value

discriminator = keras.Sequential(
    [
        keras.Input(shape=(28, 28, 1)),
        keras.layers.Conv2D(64, (3, 3), strides=(2, 2), padding="same"),
        keras.layers.LeakyReLU(negative_slope=0.2),
        keras.layers.Conv2D(128, (3, 3), strides=(2, 2), padding="same"),
        keras.layers.LeakyReLU(negative_slope=0.2),
        keras.layers.GlobalMaxPooling2D(),
        keras.layers.Dense(1),
    ],
    name="discriminator",
)
discriminator.summary()

latent_dim = 128

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

# 생성자와 생성자를 위한 옵티마이저 각각을 인스턴스화합니다.
d_optimizer = keras.optimizers.Adam(learning_rate=0.0003)
g_optimizer = keras.optimizers.Adam(learning_rate=0.0004)

# 손실 함수를 인스턴스화합니다.
loss_fn = keras.losses.BinaryCrossentropy(from_logits=True)


@tf.function
def train_step(real_images):
    # 잠재 공간에서 랜덤한 포인트를 샘플링합니다.
    random_latent_vectors = tf.random.normal(shape=(batch_size, latent_dim))
    # 샘플을 가짜 이미지로 디코딩합니다.
    generated_images = generator(random_latent_vectors)
    # 가짜 이미지와 실제 이미지를 결합합니다.
    combined_images = tf.concat([generated_images, real_images], axis=0)

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

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

    # 잠재 공간에서 랜덤한 포인트를 샘플링합니다.
    random_latent_vectors = tf.random.normal(shape=(batch_size, latent_dim))
    # "모든 이미지는 진짜 (all real images)"라고 말하는 레이블을 조립합니다.
    misleading_labels = tf.zeros((batch_size, 1))

    # 생성자를 트레이닝합니다.
    # (여기서 판별자의 가중치는 *절대로* 업데이트하면 안 됩니다)!
    with tf.GradientTape() as tape:
        predictions = discriminator(generator(random_latent_vectors))
        g_loss = loss_fn(misleading_labels, predictions)
    grads = tape.gradient(g_loss, generator.trainable_weights)
    g_optimizer.apply(grads, generator.trainable_weights)
    return d_loss, g_loss, generated_images

# 데이터셋 준비. 트레이닝 및 테스트 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)

epochs = 1  # 실제로는 멋진 숫자를 생성하려면 최소 20 에포크가 필요합니다.
save_dir = "./"

for epoch in range(epochs):
    print(f"\nStart epoch {epoch}")

    for step, real_images in enumerate(dataset):
        # 한 배치의 실제 이미지에 대해 판별기 및 생성기를 트레이닝합니다.
        d_loss, g_loss, generated_images = train_step(real_images)

        # 로깅.
        if step % 100 == 0:
            # 메트릭 출력
            print(f"discriminator loss at step {step}: {d_loss:.2f}")
            print(f"adversarial loss at step {step}: {g_loss:.2f}")

            # 생성된 이미지 중 하나를 저장합니다.
            img = keras.utils.array_to_img(generated_images[0] * 255.0, scale=False)
            img.save(os.path.join(save_dir, f"generated_img_{step}.png"))

        # 실행 시간을 제한하기 위해 10 스텝 후에 중지합니다.
        # 실제로 모델을 트레이닝하려면, 아래의 줄을 제거하세요!
        if step > 10:
            break