# 파이썬 3에서 쓰던 문법을 파이썬 2에서 사용가능하게 해줌
from __future__ import absolute_import, division, print_function, unicode_literals

!pip install -q tensorflow-gpu==2.0.0-rc1  # q 옵션은 quiet 이다.. 한참찾았네..

import tensorflow as tf

from tensorflow.keras.layers import Dense, Flatten, Conv2D
from tensorflow.keras import Model

mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0


# 채널을 나타내는 차원을 추가합니다. mnist는 흑백이므로 하나만.
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]
x_train.shape

# tf.data에 있는 slices를 사용하여 섞고, batch 단위로 나누자.
train_ds = tf.data.Dataset.from_tensor_slices(
    (x_train, y_train)
).shuffle(10000).batch(32)
test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32)

# 케라스(Keras)의 모델 서브클래싱(subclassing) API를 사용하여 tf.keras 모델을 만든다.
class MyModel(Model):
  def __init__(self):
    super(MyModel, self).__init__()
    self.conv1 = Conv2D(32, 3, activation='relu')
    self.flatten = Flatten()
    self.d1 = Dense(128, activation='relu')
    self.d2 = Dense(10, activation='relu')

  def call(self, x):
    x = self.conv1(x)
    x = self.flatten(x)
    x = self.d1(x)
    return self.d2(x)


model = MyModel()

# 훈련에 사용할 optimizer와 loss function을 선택한다.

loss_object = tf.keras.losses.SparseCategoricalCrossentropy()

optimizer = tf.keras.optimizers.Adam()

# metric을 선택한다.
train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')

test_loss = tf.keras.metrics.Mean(name='test_loss')
test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')

# tf.GradientTape를 사용하여 모델을 훈련한다.
# 이것은 모든 작업을 기록한다는 의미에서 Tape를 사용한 듯하다.

@tf.function # decorator를 사용하여 진행한다.
def train_step(images, labels):
  # with 문은 자주 사용하는 함수를 클래스로 묶은 것을 가져와서 사용할 수 있게한다.
  # 재사용에 매우 유리하다.
  # 작동 방식은, GradientTape라는 클래스 안에는 시작, 끝을 알리는 행위가 기록되어 있다.
  # 우리가 하는 작업 이후에 이를 꼭 해줘야 한다면, with문을 사용할 경우 이를 무조건 보장한다.
  # 보통 파일 읽기, 쓰기와 같은 작업에서 많이 사용한다고 보면 된다.
  # 여기서는 GradientTape에서 with 문 아래에 있는 context의 내용을 실행하고,
  # 모든 연산을 "기록"한다.
  # 우리는 train data를 넣고, loss를 뽑는 작업 이외에 수행하는 것은 항상 같다.
  # 따라서 with 문으로 설계하는 것이 맞다.
  with tf.GradientTape() as tape:
    predictions = model(images) # 예측 결과를 저장한다.
    loss = loss_object(labels, predictions)

  # 기록을 완료한뒤, 기록한 gradient를 뽑자. 기록된 tape에서 후진 자동 미분을 수행한다.
  gradients = tape.gradient(loss, model.trainable_variables)


  # 계산된 gradient를 적용한다.
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))

  # 현재 loss와 accuracy를 저장한다.
  train_loss(loss)
  train_accuracy(labels, predictions)

  [
  <tf.Variable 'my_model/conv2d/kernel:0' shape=(3, 3, 1, 32) dtype=float32>,
  <tf.Variable 'my_model/conv2d/bias:0' shape=(32,) dtype=float32>,
  <tf.Variable 'my_model/dense/kernel:0' shape=(21632, 128) dtype=float32>,
  <tf.Variable 'my_model/dense/bias:0' shape=(128,) dtype=float32>,
  <tf.Variable 'my_model/dense_1/kernel:0' shape=(128, 10) dtype=float32>,
  <tf.Variable 'my_model/dense_1/bias:0' shape=(10,) dtype=float32>]

# test에서도 만들어준다.
# 이경우는 gradient를 할필요가 없다.
@tf.function
def test_step(images, labels):
  predictions = model(images)
  t_loss = loss_object(labels, predictions)

  test_loss(t_loss)
  test_accuracy(labels, predictions)

EPOCHS = 5

for epoch in range(EPOCHS):
  for images, labels in train_ds:
    train_step(images, labels)

  for test_images, test_labels in test_ds:
    test_step(test_images, test_labels)

  template = '에포크: {}, 손실: {}, 정확도: {}, 테스트 손실: {}, 테스트 정확도: {}'
  print(template.format(epoch+1,
                        train_loss.result(),
                        train_accuracy.result()*100,
                        test_loss.result(),
                        test_accuracy.result()*100))