하위 클래스화를 통한 새로운 레이어 및 모델 만들기

TensorFlow.org에서 보기 Google Colab에서 실행하기 GitHub에서소스 보기 노트북 다운로드하기

!pip install -U tf-hub-nightly
import tensorflow_hub as hub

from tensorflow.keras import layers

import tensorflow as tf
from tensorflow import keras
2022-12-14 22:54:51.448414: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2022-12-14 22:54:51.448525: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2022-12-14 22:54:51.448535: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.

Layer 클래스: 상태(가중치)와 일부 계산의 조합

Keras의 주요 추상화 중 하나는 Layer 클래스입니다. 레이어는 상태(레이어의 "가중치")와 입력에서 출력으로의 변환("호출, 레이어의 정방향 패스")을 모두 캡슐화합니다.

다음은 밀집 레이어입니다. 상태는 변수 wb입니다.

class Linear(keras.layers.Layer):
    def __init__(self, units=32, input_dim=32):
        super(Linear, self).__init__()
        w_init = tf.random_normal_initializer()
        self.w = tf.Variable(
            initial_value=w_init(shape=(input_dim, units), dtype="float32"),
            trainable=True,
        )
        b_init = tf.zeros_initializer()
        self.b = tf.Variable(
            initial_value=b_init(shape=(units,), dtype="float32"), trainable=True
        )

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

파이썬 함수와 매우 유사한 일부 텐서 입력에서 레이어를 호출하여 레이어를 사용합니다.

x = tf.ones((2, 2))
linear_layer = Linear(4, 2)
y = linear_layer(x)
print(y)
tf.Tensor(
[[-0.07012919 -0.02425323  0.06503569  0.01101883]
 [-0.07012919 -0.02425323  0.06503569  0.01101883]], shape=(2, 4), dtype=float32)

가중치 wb는 레이어 속성으로 설정될 때 레이어에 의해 자동으로 추적됩니다.

assert linear_layer.weights == [linear_layer.w, linear_layer.b]

레이어에 가중치를 추가하는 더 빠른 바로 가기에 액세스할 수도 있습니다. add_weight() 메서드:

class Linear(keras.layers.Layer):
    def __init__(self, units=32, input_dim=32):
        super(Linear, self).__init__()
        self.w = self.add_weight(
            shape=(input_dim, units), initializer="random_normal", trainable=True
        )
        self.b = self.add_weight(shape=(units,), initializer="zeros", trainable=True)

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


x = tf.ones((2, 2))
linear_layer = Linear(4, 2)
y = linear_layer(x)
print(y)
tf.Tensor(
[[ 0.01591256 -0.08179533 -0.01888242  0.05522368]
 [ 0.01591256 -0.08179533 -0.01888242  0.05522368]], shape=(2, 4), dtype=float32)

레이어는 훈련 불가능한 가중치를 가질 수 있습니다

훈련 가능한 가중치 외에도 훈련 불가능한 가중치를 레이어에 추가할 수 있습니다. 이러한 가중치는 레이어를 훈련할 때 역전파 동안 고려되지 않아야 합니다.

훈련 불가능한 가중치를 추가 및 사용하는 방법은 다음과 같습니다.

class ComputeSum(keras.layers.Layer):
    def __init__(self, input_dim):
        super(ComputeSum, self).__init__()
        self.total = tf.Variable(initial_value=tf.zeros((input_dim,)), trainable=False)

    def call(self, inputs):
        self.total.assign_add(tf.reduce_sum(inputs, axis=0))
        return self.total


x = tf.ones((2, 2))
my_sum = ComputeSum(2)
y = my_sum(x)
print(y.numpy())
y = my_sum(x)
print(y.numpy())
[2. 2.]
[4. 4.]

layer.weights의 일부이지만, 훈련 불가능한 가중치로 분류됩니다.

print("weights:", len(my_sum.weights))
print("non-trainable weights:", len(my_sum.non_trainable_weights))

# It's not included in the trainable weights:
print("trainable_weights:", my_sum.trainable_weights)
weights: 1
non-trainable weights: 1
trainable_weights: []

모범 사례: 입력 형상이 알려질 때까지 가중치 생성 지연하기

위의 Linear 레이어는 __init__()에서 가중치 wb의 형상을 계산하는 데 사용되는 input_dim 인수를 사용했습니다.

class Linear(keras.layers.Layer):
    def __init__(self, units=32, input_dim=32):
        super(Linear, self).__init__()
        self.w = self.add_weight(
            shape=(input_dim, units), initializer="random_normal", trainable=True
        )
        self.b = self.add_weight(shape=(units,), initializer="zeros", trainable=True)

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

대부분의 경우, 입력의 크기를 미리 알지 못할 수 있으며, 레이어를 인스턴스화한 후 얼마 지나지 않아 해당 값을 알게 되면 가중치를 지연 생성하고자 합니다.

Keras API에서는 레이어의 build(self, inputs_shape) 메서드에서 레이어 가중치를 만드는 것이 좋습니다. 다음과 같습니다.

class Linear(keras.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

레이어의 __call__() 메서드는 처음 호출될 때 자동으로 빌드를 실행합니다. 지연되어 사용하기 쉬운 레이어입니다.

# At instantiation, we don't know on what inputs this is going to get called
linear_layer = Linear(32)

# The layer's weights are created dynamically the first time the layer is called
y = linear_layer(x)

위에 표시된 대로 build()를 별도로 구현하면 가중치를 한 번만 생성하는 것과 모든 호출에서 가중치를 사용하는 것을 잘 구분할 수 있습니다. 그러나 일부 고급 사용자 지정 레이어의 경우 상태 생성과 계산을 분리하는 것이 비실용적일 수 있습니다. 레이어 구현자는 가중치 생성을 첫 번째 __call__()로 연기할 수 있지만 이후 호출에서 동일한 가중치를 사용하도록 주의해야 합니다. 또한 __call__()tf.function 내부에서 처음으로 실행될 가능성이 높기 때문에 __call__()에서 발생하는 모든 변수 생성은 tf.init_scope로 래핑되어야 합니다.

재귀적으로 구성 가능한 레이어

또 다른 인스턴스의 속성으로 Layer 인스턴스를 할당하면 외부 레이어가 내부 레이어로 생성한 가중치를 추적하기 시작합니다.

__init__() 메서드에서 이러한 서브 레이어를 만들고 가중치를 빌드하도록 트리거할 수 있게 첫 번째 __call__()에 그대로 두는 것이 좋습니다.

class MLPBlock(keras.layers.Layer):
    def __init__(self):
        super(MLPBlock, self).__init__()
        self.linear_1 = Linear(32)
        self.linear_2 = Linear(32)
        self.linear_3 = Linear(1)

    def call(self, inputs):
        x = self.linear_1(inputs)
        x = tf.nn.relu(x)
        x = self.linear_2(x)
        x = tf.nn.relu(x)
        return self.linear_3(x)


mlp = MLPBlock()
y = mlp(tf.ones(shape=(3, 64)))  # The first call to the `mlp` will create the weights
print("weights:", len(mlp.weights))
print("trainable weights:", len(mlp.trainable_weights))
weights: 6
trainable weights: 6

add_loss() 메서드

레이어의 call() 메서드를 작성할 때는 훈련 루프를 작성할 때 나중에 사용하려는 손실 텐서를 만들 수 있습니다. self.add_loss(value)를 호출하면 됩니다.

# A layer that creates an activity regularization loss
class ActivityRegularizationLayer(keras.layers.Layer):
    def __init__(self, rate=1e-2):
        super(ActivityRegularizationLayer, self).__init__()
        self.rate = rate

    def call(self, inputs):
        self.add_loss(self.rate * tf.reduce_sum(inputs))
        return inputs

이러한 손실(내부 레이어에서 생성된 손실 포함)은 layer.losses를 통해 검색할 수 있습니다. 이 속성은 모든 __call__()이 시작될 때 최상위 레이어로 재설정되므로 layer.losses에는 항상 마지막 정방향 패스에서 생성된 손실값이 포함됩니다.

class OuterLayer(keras.layers.Layer):
    def __init__(self):
        super(OuterLayer, self).__init__()
        self.activity_reg = ActivityRegularizationLayer(1e-2)

    def call(self, inputs):
        return self.activity_reg(inputs)


layer = OuterLayer()
assert len(layer.losses) == 0  # No losses yet since the layer has never been called

_ = layer(tf.zeros(1, 1))
assert len(layer.losses) == 1  # We created one loss value

# `layer.losses` gets reset at the start of each __call__
_ = layer(tf.zeros(1, 1))
assert len(layer.losses) == 1  # This is the loss created during the call above

또한, loss 속성에는 내부 레이어의 가중치에 대해 생성된 정규화 손실도 포함됩니다.

class OuterLayerWithKernelRegularizer(keras.layers.Layer):
    def __init__(self):
        super(OuterLayerWithKernelRegularizer, self).__init__()
        self.dense = keras.layers.Dense(
            32, kernel_regularizer=tf.keras.regularizers.l2(1e-3)
        )

    def call(self, inputs):
        return self.dense(inputs)


layer = OuterLayerWithKernelRegularizer()
_ = layer(tf.zeros((1, 1)))

# This is `1e-3 * sum(layer.dense.kernel ** 2)`,
# created by the `kernel_regularizer` above.
print(layer.losses)
[<tf.Tensor: shape=(), dtype=float32, numpy=0.0018883398>]

이러한 손실은 다음과 같이 훈련 루프를 작성할 때 고려됩니다.

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

# Iterate over the batches of a dataset.
for x_batch_train, y_batch_train in train_dataset:
  with tf.GradientTape() as tape:
    logits = layer(x_batch_train)  # Logits for this minibatch
    # Loss value for this minibatch
    loss_value = loss_fn(y_batch_train, logits)
    # Add extra losses created during this forward pass:
    loss_value += sum(model.losses)

  grads = tape.gradient(loss_value, model.trainable_weights)
  optimizer.apply_gradients(zip(grads, model.trainable_weights))

훈련 루프 작성에 대한 자세한 가이드는 처음부터 훈련 루프 작성하기 가이드를 참조하세요.

이러한 손실은 fit()에서도 완벽하게 작동합니다(손실이 있는 경우, 자동으로 합산되어 주 손실에 추가됨).

import numpy as np

inputs = keras.Input(shape=(3,))
outputs = ActivityRegularizationLayer()(inputs)
model = keras.Model(inputs, outputs)

# If there is a loss passed in `compile`, the regularization
# losses get added to it
model.compile(optimizer="adam", loss="mse")
model.fit(np.random.random((2, 3)), np.random.random((2, 3)))

# It's also possible not to pass any loss in `compile`,
# since the model already has a loss to minimize, via the `add_loss`
# call during the forward pass!
model.compile(optimizer="adam")
model.fit(np.random.random((2, 3)), np.random.random((2, 3)))
1/1 [==============================] - 0s 179ms/step - loss: 0.1996
1/1 [==============================] - 0s 34ms/step - loss: 0.0425
<keras.callbacks.History at 0x7f44e035abb0>

add_metric() 메서드

add_loss()와 마찬가지로, 레이어에는 훈련 중 수량의 이동 평균을 추적하기 위한 add_metric() 메서드도 있습니다.

다음 "로지스틱 엔드포인트" 레이어를 고려합니다. 입력 예측 및 목표치로 사용하여 add_loss()를 통해 추적하는 손실을 계산하고 add_metric()을 통해 추적하는 정확도 스칼라를 계산합니다.

class LogisticEndpoint(keras.layers.Layer):
    def __init__(self, name=None):
        super(LogisticEndpoint, self).__init__(name=name)
        self.loss_fn = keras.losses.BinaryCrossentropy(from_logits=True)
        self.accuracy_fn = keras.metrics.BinaryAccuracy()

    def call(self, targets, logits, sample_weights=None):
        # Compute the training-time loss value and add it
        # to the layer using `self.add_loss()`.
        loss = self.loss_fn(targets, logits, sample_weights)
        self.add_loss(loss)

        # Log accuracy as a metric and add it
        # to the layer using `self.add_metric()`.
        acc = self.accuracy_fn(targets, logits, sample_weights)
        self.add_metric(acc, name="accuracy")

        # Return the inference-time prediction tensor (for `.predict()`).
        return tf.nn.softmax(logits)

이러한 방식으로 추적되는 메트릭은 layer.metrics를 통해 액세스할 수 있습니다.

layer = LogisticEndpoint()

targets = tf.ones((2, 2))
logits = tf.ones((2, 2))
y = layer(targets, logits)

print("layer.metrics:", layer.metrics)
print("current accuracy value:", float(layer.metrics[0].result()))
layer.metrics: [<keras.metrics.metrics.BinaryAccuracy object at 0x7f44e02bb9d0>]
current accuracy value: 1.0

add_loss()와 마찬가지로, 이러한 메트릭은 fit()의해 추적됩니다.

inputs = keras.Input(shape=(3,), name="inputs")
targets = keras.Input(shape=(10,), name="targets")
logits = keras.layers.Dense(10)(inputs)
predictions = LogisticEndpoint(name="predictions")(logits, targets)

model = keras.Model(inputs=[inputs, targets], outputs=predictions)
model.compile(optimizer="adam")

data = {
    "inputs": np.random.random((3, 3)),
    "targets": np.random.random((3, 10)),
}
model.fit(data)
1/1 [==============================] - 1s 588ms/step - loss: 0.8806 - binary_accuracy: 0.0000e+00
<keras.callbacks.History at 0x7f44e027e6a0>

레이어에서 선택적으로 직렬화를 활성화할 수 있습니다

함수 모델의 일부로 사용자 정의 레이어를 직렬화해야 하는 경우, 선택적으로 get_config() 메서드를 구현할 수 있습니다.

class Linear(keras.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)
{'units': 64}

기본 Layer 클래스의 __init__() 메서드는 일부 키워드 인수, 특히 namedtype를 사용합니다. 이러한 인수를 __init__()의 부모 클래스에 전달하고 레이어 구성에 포함하는 것이 좋습니다.

class Linear(keras.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)
{'name': 'linear_8', 'trainable': True, 'dtype': 'float32', 'units': 64}

구성에서 레이어를 역직렬화할 때 유연성이 더 필요한 경우, from_config() 클래스 메서드를 재정의할 수도 있습니다. 다음은 from_config()의 기본 구현입니다.

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

직렬화 및 저장에 대한 자세한 내용은 모델 저장 및 직렬화 가이드를 참조하세요.

call() 메서드의 권한 있는 training 인수

일부 레이어, 특히 BatchNormalization 레이어와 Dropout 레이어는 훈련 및 추론 중에 서로 다른 동작을 갖습니다. 이러한 레이어의 경우, call() 메서드에서 training(boolean) 인수를 노출하는 것이 표준 관행입니다.

이 인수를 call()에서 노출하면 내장 훈련 및 평가 루프(예: fit())를 사용하여 훈련 및 추론에서 레이어를 올바르게 사용할 수 있습니다.

class CustomDropout(keras.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

call() 메서드의 권한 있는 mask 인수

call()에서 지원되는 다른 권한 있는 인수는 mask 인수입니다.

이 인수는 모든 Keras RNN 레이어에서 볼 수 있습니다. 마스크는 시계열 데이터를 처리할 때 특정 입력 타임스텝을 건너뛰는 데 사용되는 부울 텐서(입력의 타임스텝당 하나의 부울 값)입니다.

Keras는 이전 레이어에서 마스크가 생성될 때 이를 지원하는 레이어에 대해 올바른 mask 인수를 __call__()에 자동으로 전달합니다. 마스크 생성 레이어는 mask_zero=True 레이어와 Masking 레이어로 구성된 Embedding입니다.

마스킹 및 마스킹 지원 레이어를 작성하는 방법에 대한 자세한 내용은 "패딩 및 마스킹 이해하기" 가이드를 확인하세요.

Model 클래스

일반적으로, Layer 클래스를 사용하여 내부 계산 블록을 정의하고 Model 클래스를 사용하여 훈련할 객체인 외부 모델을 정의합니다.

예를 들어, ResNet50 모델에는 Layer를 하위 클래스화하는 여러 ResNet 블록과 전체 ResNet50 네트워크를 포괄하는 단일 Model이 있습니다.

Model 클래스는 Layer와 같은 API를 가지며, 다음과 같은 차이점이 있습니다.

  • 내장 훈련, 평가 및 예측 루프( model.fit() , model.evaluate(), model.predict())를 제공합니다.
  • model.layers 속성을 통해 내부 레이어의 목록을 노출합니다.
  • 저장 및 직렬화 API(save(), save_weights()...)를 노출합니다.

효과적으로, Layer 클래스는 문서에서 일컫는 "레이어"("컨볼루션 레이어" 또는 "되풀이 레이어"에서와 같이) 또는 "블록"("ResNet 블록" 또는 "Inception 블록"에서와 같이)에 해당합니다.

한편, Model 클래스는 문서에서 "모델"("딥 러닝 모델"에서) 또는 "네트워크"( "딥 신경망"에서)로 지칭되는 것에 해당합다.

"Layer 클래스를 사용해야 할까요? 아니면 Model 클래스를 사용해야 할까요?"라는 질문이 있다면 자문해 보세요. fit()을 호출해야 할까? save()를 호출해야 할까? 만약 그렇다면 Model를 사용하세요. 그렇지 않다면(클래스가 더 큰 시스템의 블록이거나 직접 훈련을 작성하고 코드를 저장하기 때문에) Layer를 사용하세요.

예를 들어, 위의 mini-resnet 예제를 사용하여 fit()으로 훈련하고 save_weights()로 저장할 수 있는 Model을 빌드할 수 있습니다.

class ResNet(tf.keras.Model):

    def __init__(self, num_classes=1000):
        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(filepath)

종합: 엔드 투 엔드 예제

지금까지 배운 내용은 다음과 같습니다.

  • Layer는 상태(__init__() 또는 build()) 및 일부 계산(call()에서 정의)을 캡슐화합니다.
  • 레이어를 재귀적으로 중첩하여 새롭고 더 큰 계산 블록을 만들 수 있습니다.
  • 레이어는 add_loss()add_metric()을 통해 메트릭뿐만 아니라 손실(일반적으로, 정규화 손실)을 생성 및 추적할 수 있습니다.
  • 훈련하려는 외부 컨테이너는 Model입니다. ModelLayer와 비슷하지만, 훈련 및 직렬화 유틸리티가 추가되었습니다.

이 모든 것을 엔드 투 엔드 예제에 넣어봅시다. VAE(Variational AutoEncoder)를 구현할 것이며, MNIST 숫자로 훈련할 것입니다.

VAE는 Model의 서브 클래스가 될 것이며 Layer를 하위 클래스화하는 중첩된 레이어 구성으로 빌드됩니다. 정규화 손실(KL 확산)을 제공합니다.

from tensorflow.keras import layers


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(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

MNIST에 간단한 훈련 루프를 작성해 봅시다.

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)

epochs = 2

# Iterate over epochs.
for epoch in range(epochs):
    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_weights)
        optimizer.apply_gradients(zip(grads, vae.trainable_weights))

        loss_metric(loss)

        if step % 100 == 0:
            print("step %d: mean loss = %.4f" % (step, loss_metric.result()))
Start of epoch 0
WARNING:tensorflow:5 out of the last 5 calls to <function _BaseOptimizer._update_step_xla at 0x7f44e01a08b0> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
WARNING:tensorflow:6 out of the last 6 calls to <function _BaseOptimizer._update_step_xla at 0x7f44e01a08b0> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
step 0: mean loss = 0.3124
step 100: mean loss = 0.1235
step 200: mean loss = 0.0981
step 300: mean loss = 0.0885
step 400: mean loss = 0.0837
step 500: mean loss = 0.0805
step 600: mean loss = 0.0784
step 700: mean loss = 0.0768
step 800: mean loss = 0.0757
step 900: mean loss = 0.0747
Start of epoch 1
step 0: mean loss = 0.0744
step 100: mean loss = 0.0738
step 200: mean loss = 0.0733
step 300: mean loss = 0.0729
step 400: mean loss = 0.0725
step 500: mean loss = 0.0722
step 600: mean loss = 0.0719
step 700: mean loss = 0.0716
step 800: mean loss = 0.0714
step 900: mean loss = 0.0711

VAE는 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=2, batch_size=64)
Epoch 1/2
938/938 [==============================] - 4s 3ms/step - loss: 0.0749
Epoch 2/2
938/938 [==============================] - 3s 3ms/step - loss: 0.0677
<keras.callbacks.History at 0x7f44e005edf0>

객체 지향 개발을 넘어: 함수형 API

이 예제가 너무 지나친 객체 지향 개발입니까? 함수형 API(Functional API)를 사용하여 모델을 빌드할 수도 있습니다. 중요한 것은 하나의 스타일을 선택한다고 해서 다른 스타일로 작성된 구성 요소를 활용하지 못하는 것은 아닙니다. 항상 목적에 따라 다르게 선택할 수 있습니다.

예를 들어, 아래의 함수형 API 예제는 위 예제에서 정의한 것과 같은 Sampling 레이어를 재사용합니다.

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)
Epoch 1/3
938/938 [==============================] - 4s 3ms/step - loss: 0.0748
Epoch 2/3
938/938 [==============================] - 3s 3ms/step - loss: 0.0676
Epoch 3/3
938/938 [==============================] - 3s 3ms/step - loss: 0.0676
<keras.callbacks.History at 0x7f44681aa6a0>

자세한 정보는 함수형 API 가이드를 참고하세요.