tensorflow 2.0 keras Writing layers and models with tf keras (2)

이글은 다음 문서를 참조합니다.

www.tensorflow.org/guide/keras/custom_layers_and_models

(번역은 자력 + 파파고 + 구글 번역기를 사용하였으니, 부자연스럽더라도 양해바랍니다.)

---

You can optionally enable serialization on your layers

Functional model의 기능 중 하나로서 get_config를 사용하여 layer의 정보를 serialization할 수 있습니다. (serialization은 unit과 같은 정보를 송수신 할 수 있게 하는 것입니다.)

class Linear(layers.Layer):

  def __init__(self, units=32):
    super(Linear, self).__init__()
    self.units = units

  def build(self, input_shape):
    self.w = self.add_weight(shape=(input_shape[-1], self.units),
                             initializer='random_normal',
                             trainable=True)
    self.b = self.add_weight(shape=(self.units,),
                             initializer='random_normal',
                             trainable=True)

  def call(self, inputs):
    return tf.matmul(inputs, self.w) + self.b

  def get_config(self):
    return {'units': self.units}


# Now you can recreate the layer from its config:
layer = Linear(64)
config = layer.get_config()
print(config)
new_layer = Linear.from_config(config)

base layer class의 __init__ method가 name, dtype과 같은 keyword arguments를 가지고 있습니다. 이러한 것을 __init__의 상위 클래스에 전달하고 layer 구성에 포함시키는 것이 좋습니다.

class Linear(layers.Layer):

  def __init__(self, units=32, **kwargs):
    super(Linear, self).__init__(**kwargs)
    self.units = units

  def build(self, input_shape):
    self.w = self.add_weight(shape=(input_shape[-1], self.units),
                             initializer='random_normal',
                             trainable=True)
    self.b = self.add_weight(shape=(self.units,),
                             initializer='random_normal',
                             trainable=True)

  def call(self, inputs):
    return tf.matmul(inputs, self.w) + self.b

  def get_config(self):
    config = super(Linear, self).get_config()
    config.update({'units': self.units})
    return config


layer = Linear(64)
config = layer.get_config()
print(config)
new_layer = Linear.from_config(config)

이러한 config에서 deserializing을 하고 싶다면 from_config를 이용하면 좋습니다.

def from_config(cls, config):
  return cls(**config)

이 내용은 다음 글에서 더 자세히 다룹니다.

 

Privileged training argument in the call method

몇몇 BatchNormalization, Dropout과 같은 layer는 학습, 추론을 하는 동안 다른 기능을 가집니다. 이러한 layer의 경우 학습에 대한 여부를 boolean의 형태로 노출시키는 것이 표준 관례입니다.

이러한 인자는 fit()에서 사용하기 더 쉽습니다.

class CustomDropout(layers.Layer):

  def __init__(self, rate, **kwargs):
    super(CustomDropout, self).__init__(**kwargs)
    self.rate = rate

  def call(self, inputs, training=None):
    if training:
        return tf.nn.dropout(inputs, rate=self.rate)
    return inputs

 

Building Models

The Model class

일반적으로 block을 쌓기 위해 Layer 클래스를 사용하고, 학습을 위해 Model 객체를 정의합니다.

예를 들어 ResNet50에서는 여러 Layer로 subclassing된 ResNet block이 존재하고 Model에 통과시킵니다.

Model 클래스는 Layer와 같은 API를 가지고 있으며 다음과 같은 차이점을 가집니다.

• model.fit(), model.evaluate(), model.predict()를 가집니다. 
• model.layers 속성을 통해 내부 계층을 리스트로 접근할 수 있습니다. 
• saving and serialization API를 가집니다.

사실 Layer 클래스는 "layer"(conv layer or recurrent layer)과 "block"(ResNet block or Inception block)과 의미를 공유합니다.

한편, Model 클래스는 문맥적으로 "model"(deep learning model) 과 "network"(deep neural network)와 의미를 공유합니다.

Model을 정의하여 fit()을 통해 학습하고, save_weights를 통해 weight를 저장할 수 있습니다.

class ResNet(tf.keras.Model):

    def __init__(self):
        super(ResNet, self).__init__()
        self.block_1 = ResNetBlock()
        self.block_2 = ResNetBlock()
        self.global_pool = layers.GlobalAveragePooling2D()
        self.classifier = Dense(num_classes)

    def call(self, inputs):
        x = self.block_1(inputs)
        x = self.block_2(x)
        x = self.global_pool(x)
        return self.classifier(x)


resnet = ResNet()
dataset = ...
resnet.fit(dataset, epochs=10)
resnet.save_weights(filepath)

 

Putting it all together: an end-to-end example

지금까지 다음과 같은 내용을 보았습니다.

• Layer는 __init__ or build에서 만들어진 상태를 캡슐화하고 call을 통한 계산을 행합니다. 
• 더 새롭고 더 큰 계산 블록으로 만들어질 수 있습니다.
• (regularization)losses에 접근할 수 있습니다.
• 외부 컨테이너에 해당하는 Model로 학습시킬 수 있습니다.
• Model은 Layer을 override하고 있으나 학습과 serialization 측면에서 차이점을 가집니다.

여태까지 배운 모든 내용을 포함한 예제를 보겠습니다. MNIST를 학습시키는 Variational AutoEncoder예제입니다. VAE는 하위 클래스 계층의 중첩된 구성으로 구축된 모델의 하위 클래스입니다. 또한, regularization loss(KL divergence)에 대해서도 다룹니다.

(개념에 대한 부분은 넘어가도 좋습니다. VAE의 개념은 처음에 이해하기엔 어렵습니다. 하지만 코드가 어떤 구조를 가지고 있는지는 관찰하고 넘어가시는 것을 추천드립니다.)

class Sampling(layers.Layer):
  """Uses (z_mean, z_log_var) to sample z, the vector encoding a digit."""

  def call(self, inputs):
    z_mean, z_log_var = inputs
    batch = tf.shape(z_mean)[0]
    dim = tf.shape(z_mean)[1]
    epsilon = tf.keras.backend.random_normal(shape=(batch, dim))
    return z_mean + tf.exp(0.5 * z_log_var) * epsilon


class Encoder(layers.Layer):
  """Maps MNIST digits to a triplet (z_mean, z_log_var, z)."""

  def __init__(self,
               latent_dim=32,
               intermediate_dim=64,
               name='encoder',
               **kwargs):
    super(Encoder, self).__init__(name=name, **kwargs)
    self.dense_proj = layers.Dense(intermediate_dim, activation='relu')
    self.dense_mean = layers.Dense(latent_dim)
    self.dense_log_var = layers.Dense(latent_dim)
    self.sampling = Sampling()

  def call(self, inputs):
    x = self.dense_proj(inputs)
    z_mean = self.dense_mean(x)
    z_log_var = self.dense_log_var(x)
    z = self.sampling((z_mean, z_log_var))
    return z_mean, z_log_var, z


class Decoder(layers.Layer):
  """Converts z, the encoded digit vector, back into a readable digit."""

  def __init__(self,
               original_dim,
               intermediate_dim=64,
               name='decoder',
               **kwargs):
    super(Decoder, self).__init__(name=name, **kwargs)
    self.dense_proj = layers.Dense(intermediate_dim, activation='relu')
    self.dense_output = layers.Dense(original_dim, activation='sigmoid')

  def call(self, inputs):
    x = self.dense_proj(inputs)
    return self.dense_output(x)


class VariationalAutoEncoder(tf.keras.Model):
  """Combines the encoder and decoder into an end-to-end model for training."""

  def __init__(self,
               original_dim,
               intermediate_dim=64,
               latent_dim=32,
               name='autoencoder',
               **kwargs):
    super(VariationalAutoEncoder, self).__init__(name=name, **kwargs)
    self.original_dim = original_dim
    self.encoder = Encoder(latent_dim=latent_dim,
                           intermediate_dim=intermediate_dim)
    self.decoder = Decoder(original_dim, intermediate_dim=intermediate_dim)

  def call(self, inputs):
    z_mean, z_log_var, z = self.encoder(inputs)
    reconstructed = self.decoder(z)
    # Add KL divergence regularization loss.
    kl_loss = - 0.5 * tf.reduce_mean(
        z_log_var - tf.square(z_mean) - tf.exp(z_log_var) + 1)
    self.add_loss(kl_loss)
    return reconstructed


original_dim = 784
vae = VariationalAutoEncoder(original_dim, 64, 32)

optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
mse_loss_fn = tf.keras.losses.MeanSquaredError()

loss_metric = tf.keras.metrics.Mean()

(x_train, _), _ = tf.keras.datasets.mnist.load_data()
x_train = x_train.reshape(60000, 784).astype('float32') / 255

train_dataset = tf.data.Dataset.from_tensor_slices(x_train)
train_dataset = train_dataset.shuffle(buffer_size=1024).batch(64)

# Iterate over epochs.
for epoch in range(3):
  print('Start of epoch %d' % (epoch,))

  # Iterate over the batches of the dataset.
  for step, x_batch_train in enumerate(train_dataset):
    with tf.GradientTape() as tape:
      reconstructed = vae(x_batch_train)
      # Compute reconstruction loss
      loss = mse_loss_fn(x_batch_train, reconstructed)
      loss += sum(vae.losses)  # Add KLD regularization loss

    grads = tape.gradient(loss, vae.trainable_variables)
    optimizer.apply_gradients(zip(grads, vae.trainable_variables))

    loss_metric(loss)

    if step % 100 == 0:
      print('step %s: mean loss = %s' % (step, loss_metric.result()))

VAE는 Model의 하위클래스이기 때문에 다음과 같이 학습시킬수도 있습니다.
(데이터가 model에 어떻게 흘러가는지도 잘 관찰하는게 좋습니다)

vae = VariationalAutoEncoder(784, 64, 32)

optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)

vae.compile(optimizer, loss=tf.keras.losses.MeanSquaredError())
vae.fit(x_train, x_train, epochs=3, batch_size=64)

 

Beyond object-oriented development: the Functional API

Functional API를 사용해서 구현할 수도 있습니다. 어떤 방식을 쓰던 문제가 되지 않습니다. 섞어 사용하여도 무방합니다.

original_dim = 784
intermediate_dim = 64
latent_dim = 32

# Define encoder model.
original_inputs = tf.keras.Input(shape=(original_dim,), name='encoder_input')
x = layers.Dense(intermediate_dim, activation='relu')(original_inputs)
z_mean = layers.Dense(latent_dim, name='z_mean')(x)
z_log_var = layers.Dense(latent_dim, name='z_log_var')(x)
z = Sampling()((z_mean, z_log_var))
encoder = tf.keras.Model(inputs=original_inputs, outputs=z, name='encoder')

# Define decoder model.
latent_inputs = tf.keras.Input(shape=(latent_dim,), name='z_sampling')
x = layers.Dense(intermediate_dim, activation='relu')(latent_inputs)
outputs = layers.Dense(original_dim, activation='sigmoid')(x)
decoder = tf.keras.Model(inputs=latent_inputs, outputs=outputs, name='decoder')

# Define VAE model.
outputs = decoder(z)
vae = tf.keras.Model(inputs=original_inputs, outputs=outputs, name='vae')

# Add KL divergence regularization loss.
kl_loss = - 0.5 * tf.reduce_mean(
    z_log_var - tf.square(z_mean) - tf.exp(z_log_var) + 1)
vae.add_loss(kl_loss)

# Train.
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
vae.compile(optimizer, loss=tf.keras.losses.MeanSquaredError())
vae.fit(x_train, x_train, epochs=3, batch_size=64)