#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Explanation: Copyright 2020 The TensorFlow Authors.
End of explanation
"""


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

"""
Explanation: 처음부터 훈련 루프 작성하기
<table class="tfo-notebook-buttons" align="left">
  <td><a target="_blank" href="https://www.tensorflow.org/guide/keras/writing_a_training_loop_from_scratch"><img src="https://www.tensorflow.org/images/tf_logo_32px.png">TensorFlow.org에서 보기</a></td>
  <td><a target="_blank" href="https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/ko/guide/keras/writing_a_training_loop_from_scratch.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png">Google Colab에서 실행</a></td>
  <td><a target="_blank" href="https://github.com/tensorflow/docs-l10n/blob/master/site/ko/guide/keras/writing_a_training_loop_from_scratch.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png"> TensorFlow.org에서 보기</a></td>
  <td><a href="https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/ko/guide/keras/writing_a_training_loop_from_scratch.ipynb"><img src="https://www.tensorflow.org/images/download_logo_32px.png">노트북 다운로드</a></td>
</table>

Setup
End of explanation
"""


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

"""
Explanation: 시작하기
Keras는 기본 학습 및 평가 루프인 fit() 및 evaluate()를 제공합니다. 사용법은 내장된 메서드를 이용한 훈련 및 평가 가이드에서 다룹니다.
fit()의 편리성을 이용하면서 해당 모델의 학습 알고리즘을 사용자 정의하려면(예를 들어, fit()를 사용하여 GAN을 훈련시킴), Model 클래스를 하위 클래스화하고 fit() 중에 반복적으로 호출되는 고유 train_step() 메서드를 구현할 수 있습니다. 이 내용은 <a data-md-type="raw_html" href="https://www.tensorflow.org/guide/keras/customizing_what_happens_in_fit/">fit() 동작 사용자 정의하기</a> 가이드에서 다룹니다.
이제 교육 및 평가에 대한 매우 낮은 수준의 제어를 원한다면 자체 교육 및 평가 루프를 처음부터 작성해야합니다. 이것이이 안내서의 내용입니다.
GradientTape 사용하기: 첫 번째 엔드 투 엔드 예제
GradientTape 범위 내에서 모델을 호출하면 손실 값과 관련하여 레이어의 훈련 가능한 가중치 기울기를 가져올 수 있습니다. 최적화 프로그램 인스턴스를 사용하면 이러한 기울기를 이용해 이러한 변수를 업데이트할 수 있습니다(model.trainable_weights를 사용하여 가져올 수 있음).
간단한 MNIST 모델을 생각해 보겠습니다.
End of explanation
"""


# Instantiate an optimizer.
optimizer = keras.optimizers.SGD(learning_rate=1e-3)
# Instantiate a loss function.
loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# Prepare the training dataset.
batch_size = 64
(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))

# Reserve 10,000 samples for validation.
x_val = x_train[-10000:]
y_val = y_train[-10000:]
x_train = x_train[:-10000]
y_train = y_train[:-10000]

# Prepare the training dataset.
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = train_dataset.shuffle(buffer_size=1024).batch(batch_size)

# Prepare the validation dataset.
val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))
val_dataset = val_dataset.batch(batch_size)

"""
Explanation: 사용자 정의 훈련 루프와 함께 미니 배치 기울기를 사용해 이를 훈련시키겠습니다.
우선 최적화 프로그램, 손실 함수 및 데이터세트가 필요합니다.
End of explanation
"""


epochs = 2
for epoch in range(epochs):
    print("\nStart of epoch %d" % (epoch,))

    # Iterate over the batches of the dataset.
    for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):

        # Open a GradientTape to record the operations run
        # during the forward pass, which enables auto-differentiation.
        with tf.GradientTape() as tape:

            # Run the forward pass of the layer.
            # The operations that the layer applies
            # to its inputs are going to be recorded
            # on the GradientTape.
            logits = model(x_batch_train, training=True)  # Logits for this minibatch

            # Compute the loss value for this minibatch.
            loss_value = loss_fn(y_batch_train, logits)

        # Use the gradient tape to automatically retrieve
        # the gradients of the trainable variables with respect to the loss.
        grads = tape.gradient(loss_value, model.trainable_weights)

        # Run one step of gradient descent by updating
        # the value of the variables to minimize the loss.
        optimizer.apply_gradients(zip(grads, model.trainable_weights))

        # Log every 200 batches.
        if step % 200 == 0:
            print(
                "Training loss (for one batch) at step %d: %.4f"
                % (step, float(loss_value))
            )
            print("Seen so far: %s samples" % ((step + 1) * batch_size))

"""
Explanation: 사용자 정의 훈련 루프

우리는 시대를 반복하는 for 루프를 엽니 다.
각 시대마다 데이터 집합을 반복적으로 처리하는 for 루프를 엽니 다.
각 배치마다 GradientTape() 범위를 엽니 다
이 범위 내에서 모델을 호출하고(순방향 전달) 손실을 계산합니다.
스코프 외부에서 손실과 관련하여 모델 가중치의 기울기를 검색합니다.
마지막으로, 그라디언트를 기반으로 모델의 가중치를 업데이트하기 위해 옵티 마이저를 사용합니다
End of explanation
"""


# Get model
inputs = keras.Input(shape=(784,), name="digits")
x = layers.Dense(64, activation="relu", name="dense_1")(inputs)
x = layers.Dense(64, activation="relu", name="dense_2")(x)
outputs = layers.Dense(10, name="predictions")(x)
model = keras.Model(inputs=inputs, outputs=outputs)

# Instantiate an optimizer to train the model.
optimizer = keras.optimizers.SGD(learning_rate=1e-3)
# Instantiate a loss function.
loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# Prepare the metrics.
train_acc_metric = keras.metrics.SparseCategoricalAccuracy()
val_acc_metric = keras.metrics.SparseCategoricalAccuracy()

"""
Explanation: 낮은 수준의 메트릭 처리
이 기본 루프에 메트릭 모니터링을 추가하겠습니다.
처음부터 작성한 이러한 훈련 루프에서 내장 메트릭(또는 사용자가 작성한 사용자 정의 메트릭)을 쉽게 재사용할 수 있습니다. 흐름은 다음과 같습니다.

루프가 시작되는 부분에 메트릭 인스턴스화
각 배치 끝에서 metric.update_state() 호출
메트릭의 현재 값을 표시해야 할 때 metric.result() 호출
메트릭의 상태를 지워야 할 때(일반적으로 epoch의 끝에서) metric.reset_states() 호출

이러한 지식을 바탕으로 각 epoch가 끝날 때 유효성 검사 데이터에 대해 SparseCategoricalAccuracy를 계산하겠습니다.
End of explanation
"""


import time

epochs = 2
for epoch in range(epochs):
    print("\nStart of epoch %d" % (epoch,))
    start_time = time.time()

    # Iterate over the batches of the dataset.
    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_gradients(zip(grads, model.trainable_weights))

        # Update training metric.
        train_acc_metric.update_state(y_batch_train, logits)

        # Log every 200 batches.
        if step % 200 == 0:
            print(
                "Training loss (for one batch) at step %d: %.4f"
                % (step, float(loss_value))
            )
            print("Seen so far: %d samples" % ((step + 1) * batch_size))

    # Display metrics at the end of each epoch.
    train_acc = train_acc_metric.result()
    print("Training acc over epoch: %.4f" % (float(train_acc),))

    # Reset training metrics at the end of each epoch
    train_acc_metric.reset_states()

    # Run a validation loop at the end of each epoch.
    for x_batch_val, y_batch_val in val_dataset:
        val_logits = model(x_batch_val, training=False)
        # Update val metrics
        val_acc_metric.update_state(y_batch_val, val_logits)
    val_acc = val_acc_metric.result()
    val_acc_metric.reset_states()
    print("Validation acc: %.4f" % (float(val_acc),))
    print("Time taken: %.2fs" % (time.time() - start_time))

"""
Explanation: 훈련 및 평가 루프는 다음과 같습니다.
End of explanation
"""


@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_gradients(zip(grads, model.trainable_weights))
    train_acc_metric.update_state(y, logits)
    return loss_value


"""
Explanation: tf.function을 이용해 훈련 단계의 속도 높이기
TensorFlow 2.0의 기본 런타임은 즉시 실행형입니다. 따라서 위의 훈련 루프는 즉시 실행됩니다.
이것은 디버깅에는 유용하지만 그래프 컴파일에는 확실한 성능 이점이 있습니다. 계산을 정적 그래프로 정의하면 프레임 워크에서 전체 성능 최적화를 적용 할 수 있습니다. 프레임 워크가 다음 작업을 알지 못하고 탐욕스럽게 하나의 작업을 실행하도록 제한되어 있으면 불가능합니다.
텐서를 입력으로 사용하는 모든 함수를 정적 그래프로 컴파일 할 수 있습니다. 다음과 같이 @tf.function 데코레이터를 추가하십시오.
End of explanation
"""


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


"""
Explanation: 평가 단계에도 동일한 작업을 수행합니다.
End of explanation
"""


import time

epochs = 2
for epoch in range(epochs):
    print("\nStart of epoch %d" % (epoch,))
    start_time = time.time()

    # Iterate over the batches of the dataset.
    for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):
        loss_value = train_step(x_batch_train, y_batch_train)

        # Log every 200 batches.
        if step % 200 == 0:
            print(
                "Training loss (for one batch) at step %d: %.4f"
                % (step, float(loss_value))
            )
            print("Seen so far: %d samples" % ((step + 1) * batch_size))

    # Display metrics at the end of each epoch.
    train_acc = train_acc_metric.result()
    print("Training acc over epoch: %.4f" % (float(train_acc),))

    # Reset training metrics at the end of each epoch
    train_acc_metric.reset_states()

    # Run a validation loop at the end of each epoch.
    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_states()
    print("Validation acc: %.4f" % (float(val_acc),))
    print("Time taken: %.2fs" % (time.time() - start_time))

"""
Explanation: 이제 이 컴파일된 훈련 단계를 이용해 훈련 루프를 다시 실행해 보겠습니다.
End of explanation
"""


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


"""
Explanation: 훨씬 빠르지 않나요?
모델에 의해 추적되는 손실의 저수준 처리
레이어 및 모델은 self.add_loss(value) 를 호출하는 레이어에 의해 순방향 전달 중에 생성되는 모든 손실을 재귀적으로 추적합니다. 스칼라 손실 값의 결과 목록은 순방향 전달이 끝날 때 model.losses 속성을 통해 이용할 수 있습니다.
이러한 손실 구성 요소를 사용하려면 해당 구성 요소를 합산하여 교육 단계의 주요 손실에 추가해야합니다.
이 계층을 고려하면 활동 정규화 손실이 발생합니다.
End of explanation
"""


inputs = keras.Input(shape=(784,), name="digits")
x = layers.Dense(64, activation="relu")(inputs)
# Insert activity regularization as a layer
x = ActivityRegularizationLayer()(x)
x = layers.Dense(64, activation="relu")(x)
outputs = layers.Dense(10, name="predictions")(x)

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

"""
Explanation: 이것을 사용하는 정말 간단한 모델을 만들어 보겠습니다.
End of explanation
"""


@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        logits = model(x, training=True)
        loss_value = loss_fn(y, logits)
        # Add any extra losses created during the forward pass.
        loss_value += sum(model.losses)
    grads = tape.gradient(loss_value, model.trainable_weights)
    optimizer.apply_gradients(zip(grads, model.trainable_weights))
    train_acc_metric.update_state(y, logits)
    return loss_value


"""
Explanation: 훈련 단계는 다음과 같습니다.
End of explanation
"""


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

"""
Explanation: 요약
이제 내장된 훈련 루프를 사용하고 처음부터 고유한 루프를 작성하는 방법에 대해 알아야 할 모든 내용을 살펴보았습니다.
결론적으로, 다음은이 가이드에서 배운 모든 것, 즉 MNIST 숫자에 대해 훈련 된 DCGAN을 하나로 묶는 간단한 종단 간 예제입니다.
엔드-투-엔드 예제: 처음부터 GAN 훈련 루프 만들기
아마도 GAN(Generative Adversarial Networks)에 익숙할 것으로 생각합니다. GAN은 이미지의 훈련 데이터세트의 잠재 분포(이미지의 "잠재적 공간")를 학습함으로써 거의 실제처럼 보이는 새로운 이미지를 생성할 수 있습니다.
GAN은 잠재적 공간의 포인트를 이미지 공간의 포인트에 매핑하는 "생성기(generator)" 모델과 실제 이미지(훈련 데이터세트로 만들어짐)와 가짜 이미지(생성기 네트워크의 출력)의 차이를 구분할 수 있는 분류기인 "판별자(discriminator)" 모델의 두 부분으로 이루어집니다.
A GAN training loop looks like this:


차별자를 훈련 시키십시오.


잠재 공간에서 임의의 지점을 일괄 적으로 샘플링합니다.

"생성기" 모델을 통해 포인트를 가짜 이미지로 바꿉니다.
실제 이미지 배치를 가져와 생성된 이미지와 결합합니다.

생성 된 이미지와 실제 이미지를 분류하기 위해 "차별 자"모델을 훈련시킵니다.


생성기를 훈련시킵니다.


잠재 공간에서 임의의 포인트를 샘플링합니다.

"발전기"네트워크를 통해 포인트를 가짜 이미지로 바꾸십시오.
실제 이미지 배치를 가져와 생성된 이미지와 결합합니다.
"생성기" 모델을 훈련시켜 판별자를 "속이고" 가짜 이미지를 실제 이미지로 분류합니다.

GAN 작동 방식에 대한 자세한 개요는 Python을 사용한 딥 러닝을 참조하세요.
이 훈련 루프를 구현해 보겠습니다. 먼저 가짜 숫자와 진짜 숫자를 분류하도록 의도된 판별자를 만듭니다.
End of explanation
"""


latent_dim = 128

generator = keras.Sequential(
    [
        keras.Input(shape=(latent_dim,)),
        # We want to generate 128 coefficients to reshape into a 7x7x128 map
        layers.Dense(7 * 7 * 128),
        layers.LeakyReLU(alpha=0.2),
        layers.Reshape((7, 7, 128)),
        layers.Conv2DTranspose(128, (4, 4), strides=(2, 2), padding="same"),
        layers.LeakyReLU(alpha=0.2),
        layers.Conv2DTranspose(128, (4, 4), strides=(2, 2), padding="same"),
        layers.LeakyReLU(alpha=0.2),
        layers.Conv2D(1, (7, 7), padding="same", activation="sigmoid"),
    ],
    name="generator",
)

"""
Explanation: 그런 다음, 잠재 벡터를 형상 (28, 28, 1)(MNIST 숫자를 나타냄)의 출력으로 바꾸는 생성기 네트워크를 만듭니다.
End of explanation
"""


# Instantiate one optimizer for the discriminator and another for the generator.
d_optimizer = keras.optimizers.Adam(learning_rate=0.0003)
g_optimizer = keras.optimizers.Adam(learning_rate=0.0004)

# Instantiate a loss function.
loss_fn = keras.losses.BinaryCrossentropy(from_logits=True)


@tf.function
def train_step(real_images):
    # Sample random points in the latent space
    random_latent_vectors = tf.random.normal(shape=(batch_size, latent_dim))
    # Decode them to fake images
    generated_images = generator(random_latent_vectors)
    # Combine them with real images
    combined_images = tf.concat([generated_images, real_images], axis=0)

    # Assemble labels discriminating real from fake images
    labels = tf.concat(
        [tf.ones((batch_size, 1)), tf.zeros((real_images.shape[0], 1))], axis=0
    )
    # Add random noise to the labels - important trick!
    labels += 0.05 * tf.random.uniform(labels.shape)

    # Train the discriminator
    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_gradients(zip(grads, discriminator.trainable_weights))

    # Sample random points in the latent space
    random_latent_vectors = tf.random.normal(shape=(batch_size, latent_dim))
    # Assemble labels that say "all real images"
    misleading_labels = tf.zeros((batch_size, 1))

    # Train the generator (note that we should *not* update the weights
    # of the discriminator)!
    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_gradients(zip(grads, generator.trainable_weights))
    return d_loss, g_loss, generated_images


"""
Explanation: 핵심 비트는 훈련 루프입니다. 보시다시피 매우 간단합니다. 훈련 단계 기능은 17 줄만 사용합니다.
End of explanation
"""


import os

# Prepare the dataset. We use both the training & test MNIST digits.
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  # In practice you need at least 20 epochs to generate nice digits.
save_dir = "./"

for epoch in range(epochs):
    print("\nStart epoch", epoch)

    for step, real_images in enumerate(dataset):
        # Train the discriminator & generator on one batch of real images.
        d_loss, g_loss, generated_images = train_step(real_images)

        # Logging.
        if step % 200 == 0:
            # Print metrics
            print("discriminator loss at step %d: %.2f" % (step, d_loss))
            print("adversarial loss at step %d: %.2f" % (step, g_loss))

            # Save one generated image
            img = tf.keras.preprocessing.image.array_to_img(
                generated_images[0] * 255.0, scale=False
            )
            img.save(os.path.join(save_dir, "generated_img" + str(step) + ".png"))

        # To limit execution time we stop after 10 steps.
        # Remove the lines below to actually train the model!
        if step > 10:
            break

"""
Explanation: 일련의 이미지에 대해 train_step 을 반복해서 호출하여 GAN을 훈련시켜 봅시다.
판별 기 및 생성기는 convnet이므로이 코드를 GPU에서 실행하려고합니다.
End of explanation
"""


