Lưu ngày! Google I / O hoạt động trở lại từ ngày 18 đến 20 tháng 5 Đăng ký ngay
Trang này được dịch bởi Cloud Translation API.
Switch to English

Bộ mã tự động biến thể phù hợp

Xem trên TensorFlow.org Chạy trong Google Colab Xem nguồn trên GitHub Tải xuống sổ ghi chép

Sổ tay này trình bày cách đào tạo một Bộ mã tự động biến thể (VAE) ( 1 , 2 ) trên tập dữ liệu MNIST. VAE là một mô hình xác suất sử dụng tự động mã hóa, một mô hình lấy dữ liệu đầu vào có chiều cao và nén nó thành một biểu diễn nhỏ hơn. Không giống như một bộ mã tự động truyền thống, ánh xạ đầu vào vào một vectơ tiềm ẩn, một VAE ánh xạ dữ liệu đầu vào thành các tham số của phân phối xác suất, chẳng hạn như giá trị trung bình và phương sai của một Gaussian. Cách tiếp cận này tạo ra một không gian tiềm ẩn có cấu trúc, liên tục, rất hữu ích cho việc tạo hình ảnh.

Không gian tiềm ẩn của hình ảnh CVAE

Thiết lập

pip install -q tensorflow-probability

# to generate gifs
pip install -q imageio
pip install -q git+https://github.com/tensorflow/docs
from IPython import display

import glob
import imageio
import matplotlib.pyplot as plt
import numpy as np
import PIL
import tensorflow as tf
import tensorflow_probability as tfp
import time

Tải tập dữ liệu MNIST

Mỗi hình ảnh MNIST ban đầu là một vectơ gồm 784 số nguyên, mỗi số trong số đó nằm trong khoảng từ 0-255 và đại diện cho cường độ của một pixel. Lập mô hình từng pixel với phân phối Bernoulli trong mô hình của chúng tôi và mã hóa tĩnh bộ dữ liệu.

(train_images, _), (test_images, _) = tf.keras.datasets.mnist.load_data()
Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz
11493376/11490434 [==============================] - 0s 0us/step
def preprocess_images(images):
  images = images.reshape((images.shape[0], 28, 28, 1)) / 255.
  return np.where(images > .5, 1.0, 0.0).astype('float32')

train_images = preprocess_images(train_images)
test_images = preprocess_images(test_images)
train_size = 60000
batch_size = 32
test_size = 10000

Sử dụng tf.data để trộn và trộn dữ liệu

train_dataset = (tf.data.Dataset.from_tensor_slices(train_images)
                 .shuffle(train_size).batch(batch_size))
test_dataset = (tf.data.Dataset.from_tensor_slices(test_images)
                .shuffle(test_size).batch(batch_size))

Xác định bộ mã hóa và mạng giải mã với tf.keras.

Trong ví dụ VAE này, hãy sử dụng hai ConvNets nhỏ cho mạng bộ mã hóa và bộ giải mã. Trong tài liệu, những mạng lưới này còn được gọi là mô hình suy luận / công nhận và mô hình tổng hợp. Sử dụng tf.keras.Sequential để đơn giản hóa việc triển khai. Gọi $ x $ và $ z $ lần lượt biểu thị biến quan sát và biến tiềm ẩn trong các mô tả sau.

Mạng mã hóa

Điều này xác định phân phối sau gần đúng $ q (z | x) $, lấy đầu vào là một quan sát và xuất ra một tập hợp các tham số để chỉ định phân phối có điều kiện của biểu diễn tiềm ẩn $ z $. Trong ví dụ này, chỉ cần lập mô hình phân phối dưới dạng Gaussian chéo và mạng xuất ra các tham số trung bình và log-phương sai của một Gauss đã phân tích nhân tử. Xuất log-variance thay vì phương sai trực tiếp để ổn định số.

Mạng giải mã

Điều này xác định phân phối có điều kiện của quan sát $ p (x | z) $, lấy một mẫu tiềm ẩn $ z $ làm đầu vào và đầu ra các tham số cho phân phối có điều kiện của quan sát. Lập mô hình phân phối tiềm ẩn trước $ p (z) $ dưới dạng Gaussian đơn vị.

Thủ thuật phân loại lại

Để tạo một mẫu $ z $ cho bộ giải mã trong quá trình đào tạo, bạn có thể lấy mẫu từ phân phối tiềm ẩn được xác định bởi các tham số do bộ mã hóa xuất ra, với một quan sát đầu vào $ x $. Tuy nhiên, hoạt động lấy mẫu này tạo ra một nút thắt cổ chai vì sự lan truyền ngược không thể chảy qua một nút ngẫu nhiên.

Để giải quyết vấn đề này, hãy sử dụng thủ thuật đánh giá lại. Trong ví dụ của chúng tôi, bạn ước tính $ z $ bằng cách sử dụng các tham số bộ giải mã và một tham số khác $ \ epsilon $ như sau:

$$z = \mu + \sigma \odot \epsilon$$

trong đó $ \ mu $ và $ \ sigma $ đại diện cho giá trị trung bình và độ lệch chuẩn của phân phối Gaussian tương ứng. Chúng có thể được lấy từ đầu ra của bộ giải mã. $ \ Epsilon $ có thể được coi là một nhiễu ngẫu nhiên được sử dụng để duy trì ngẫu nhiên của $ z $. Tạo $ \ epsilon $ từ phân phối chuẩn chuẩn.

Biến tiềm ẩn $ z $ hiện được tạo bởi một hàm của $ \ mu $, $ \ sigma $ và $ \ epsilon $, sẽ cho phép mô hình sao chép các gradient trong bộ mã hóa thông qua $ \ mu $ và $ \ sigma $ tương ứng , đồng thời duy trì tính ngẫu nhiên thông qua $ \ epsilon $.

Kiến trúc mạng

Đối với mạng bộ mã hóa, hãy sử dụng hai lớp phức hợp theo sau là một lớp được kết nối đầy đủ. Trong mạng bộ giải mã, phản chiếu kiến ​​trúc này bằng cách sử dụng một lớp được kết nối đầy đủ theo sau là ba lớp chuyển vị tích chập (hay còn gọi là lớp giải mã trong một số ngữ cảnh). Lưu ý, thực tế phổ biến là tránh sử dụng chuẩn hóa hàng loạt khi đào tạo VAE, vì ngẫu nhiên bổ sung do sử dụng các lô nhỏ có thể làm trầm trọng thêm tính không ổn định bên cạnh ngẫu nhiên từ việc lấy mẫu.

class CVAE(tf.keras.Model):
  """Convolutional variational autoencoder."""

  def __init__(self, latent_dim):
    super(CVAE, self).__init__()
    self.latent_dim = latent_dim
    self.encoder = tf.keras.Sequential(
        [
            tf.keras.layers.InputLayer(input_shape=(28, 28, 1)),
            tf.keras.layers.Conv2D(
                filters=32, kernel_size=3, strides=(2, 2), activation='relu'),
            tf.keras.layers.Conv2D(
                filters=64, kernel_size=3, strides=(2, 2), activation='relu'),
            tf.keras.layers.Flatten(),
            # No activation
            tf.keras.layers.Dense(latent_dim + latent_dim),
        ]
    )

    self.decoder = tf.keras.Sequential(
        [
            tf.keras.layers.InputLayer(input_shape=(latent_dim,)),
            tf.keras.layers.Dense(units=7*7*32, activation=tf.nn.relu),
            tf.keras.layers.Reshape(target_shape=(7, 7, 32)),
            tf.keras.layers.Conv2DTranspose(
                filters=64, kernel_size=3, strides=2, padding='same',
                activation='relu'),
            tf.keras.layers.Conv2DTranspose(
                filters=32, kernel_size=3, strides=2, padding='same',
                activation='relu'),
            # No activation
            tf.keras.layers.Conv2DTranspose(
                filters=1, kernel_size=3, strides=1, padding='same'),
        ]
    )

  @tf.function
  def sample(self, eps=None):
    if eps is None:
      eps = tf.random.normal(shape=(100, self.latent_dim))
    return self.decode(eps, apply_sigmoid=True)

  def encode(self, x):
    mean, logvar = tf.split(self.encoder(x), num_or_size_splits=2, axis=1)
    return mean, logvar

  def reparameterize(self, mean, logvar):
    eps = tf.random.normal(shape=mean.shape)
    return eps * tf.exp(logvar * .5) + mean

  def decode(self, z, apply_sigmoid=False):
    logits = self.decoder(z)
    if apply_sigmoid:
      probs = tf.sigmoid(logits)
      return probs
    return logits

Xác định hàm mất mát và trình tối ưu hóa

VAE đào tạo bằng cách tối đa hóa giới hạn dưới bằng chứng (ELBO) về khả năng log biên:

$$\log p(x) \ge \text{ELBO} = \mathbb{E}_{q(z|x)}\left[\log \frac{p(x, z)}{q(z|x)}\right].$$

Trên thực tế, hãy tối ưu hóa ước tính Monte Carlo mẫu đơn về kỳ vọng này:

$$\log p(x| z) + \log p(z) - \log q(z|x),$$

trong đó $ z $ được lấy mẫu từ $ q (z | x) $.

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


def log_normal_pdf(sample, mean, logvar, raxis=1):
  log2pi = tf.math.log(2. * np.pi)
  return tf.reduce_sum(
      -.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi),
      axis=raxis)


def compute_loss(model, x):
  mean, logvar = model.encode(x)
  z = model.reparameterize(mean, logvar)
  x_logit = model.decode(z)
  cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(logits=x_logit, labels=x)
  logpx_z = -tf.reduce_sum(cross_ent, axis=[1, 2, 3])
  logpz = log_normal_pdf(z, 0., 0.)
  logqz_x = log_normal_pdf(z, mean, logvar)
  return -tf.reduce_mean(logpx_z + logpz - logqz_x)


@tf.function
def train_step(model, x, optimizer):
  """Executes one training step and returns the loss.

  This function computes the loss and gradients, and uses the latter to
  update the model's parameters.
  """
  with tf.GradientTape() as tape:
    loss = compute_loss(model, x)
  gradients = tape.gradient(loss, model.trainable_variables)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))

Đào tạo

  • Bắt đầu bằng cách lặp lại tập dữ liệu
  • Trong mỗi lần lặp lại, hãy chuyển hình ảnh đến bộ mã hóa để thu được một tập hợp các tham số trung bình và log-phương sai của giá trị sau gần đúng $ q (z | x) $
  • sau đó áp dụng thủ thuật phân loại lại để lấy mẫu từ $ q (z | x) $
  • Cuối cùng, chuyển các mẫu đã được phân tích lại cho bộ giải mã để nhận được log của phân phối tổng hợp $ p (x | z) $
  • Lưu ý: Vì bạn sử dụng tập dữ liệu được tải bởi keras với 60k điểm dữ liệu trong tập huấn luyện và 10k điểm dữ liệu trong tập thử nghiệm, kết quả ELBO của chúng tôi trên tập thử nghiệm cao hơn một chút so với kết quả được báo cáo trong tài liệu sử dụng binarization động MNIST của Larochelle.

Tạo hình ảnh

  • Sau khi đào tạo, đã đến lúc tạo ra một số hình ảnh
  • Bắt đầu bằng cách lấy mẫu một tập hợp các vectơ tiềm ẩn từ phân phối trước Gaussian đơn vị $ p (z) $
  • Sau đó, trình tạo sẽ chuyển đổi mẫu tiềm ẩn $ z $ thành logits của quan sát, đưa ra phân phối $ p (x | z) $
  • Ở đây, hãy vẽ biểu đồ xác suất của các phân phối Bernoulli
epochs = 10
# set the dimensionality of the latent space to a plane for visualization later
latent_dim = 2
num_examples_to_generate = 16

# keeping the random vector constant for generation (prediction) so
# it will be easier to see the improvement.
random_vector_for_generation = tf.random.normal(
    shape=[num_examples_to_generate, latent_dim])
model = CVAE(latent_dim)
def generate_and_save_images(model, epoch, test_sample):
  mean, logvar = model.encode(test_sample)
  z = model.reparameterize(mean, logvar)
  predictions = model.sample(z)
  fig = plt.figure(figsize=(4, 4))

  for i in range(predictions.shape[0]):
    plt.subplot(4, 4, i + 1)
    plt.imshow(predictions[i, :, :, 0], cmap='gray')
    plt.axis('off')

  # tight_layout minimizes the overlap between 2 sub-plots
  plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))
  plt.show()
# Pick a sample of the test set for generating output images
assert batch_size >= num_examples_to_generate
for test_batch in test_dataset.take(1):
  test_sample = test_batch[0:num_examples_to_generate, :, :, :]
generate_and_save_images(model, 0, test_sample)

for epoch in range(1, epochs + 1):
  start_time = time.time()
  for train_x in train_dataset:
    train_step(model, train_x, optimizer)
  end_time = time.time()

  loss = tf.keras.metrics.Mean()
  for test_x in test_dataset:
    loss(compute_loss(model, test_x))
  elbo = -loss.result()
  display.clear_output(wait=False)
  print('Epoch: {}, Test set ELBO: {}, time elapse for current epoch: {}'
        .format(epoch, elbo, end_time - start_time))
  generate_and_save_images(model, epoch, test_sample)
Epoch: 10, Test set ELBO: -156.66375732421875, time elapse for current epoch: 5.057319402694702

png

Hiển thị hình ảnh được tạo từ kỷ nguyên đào tạo cuối cùng

def display_image(epoch_no):
  return PIL.Image.open('image_at_epoch_{:04d}.png'.format(epoch_no))
plt.imshow(display_image(epoch))
plt.axis('off')  # Display images
(-0.5, 287.5, 287.5, -0.5)

png

Hiển thị một GIF động của tất cả các hình ảnh đã lưu

anim_file = 'cvae.gif'

with imageio.get_writer(anim_file, mode='I') as writer:
  filenames = glob.glob('image*.png')
  filenames = sorted(filenames)
  for filename in filenames:
    image = imageio.imread(filename)
    writer.append_data(image)
  image = imageio.imread(filename)
  writer.append_data(image)
import tensorflow_docs.vis.embed as embed
embed.embed_file(anim_file)

gif

Hiển thị đa dạng chữ số 2D từ không gian ẩn

Chạy mã bên dưới sẽ hiển thị sự phân bố liên tục của các lớp chữ số khác nhau, với mỗi chữ số biến đổi thành một chữ số khác trên không gian tiềm ẩn 2D. Sử dụng Xác suất TensorFlow để tạo phân phối chuẩn chuẩn cho không gian tiềm ẩn.

def plot_latent_images(model, n, digit_size=28):
  """Plots n x n digit images decoded from the latent space."""

  norm = tfp.distributions.Normal(0, 1)
  grid_x = norm.quantile(np.linspace(0.05, 0.95, n))
  grid_y = norm.quantile(np.linspace(0.05, 0.95, n))
  image_width = digit_size*n
  image_height = image_width
  image = np.zeros((image_height, image_width))

  for i, yi in enumerate(grid_x):
    for j, xi in enumerate(grid_y):
      z = np.array([[xi, yi]])
      x_decoded = model.sample(z)
      digit = tf.reshape(x_decoded[0], (digit_size, digit_size))
      image[i * digit_size: (i + 1) * digit_size,
            j * digit_size: (j + 1) * digit_size] = digit.numpy()

  plt.figure(figsize=(10, 10))
  plt.imshow(image, cmap='Greys_r')
  plt.axis('Off')
  plt.show()
plot_latent_images(model, 20)

png

Bước tiếp theo

Hướng dẫn này đã trình bày cách triển khai bộ mã tự động biến thể tích hợp bằng cách sử dụng TensorFlow.

Bước tiếp theo, bạn có thể cố gắng cải thiện đầu ra của mô hình bằng cách tăng kích thước mạng. Ví dụ: bạn có thể thử đặt các thông số filter cho từng lớp Conv2DConv2DTranspose thành 512. Lưu ý rằng để tạo biểu đồ hình ảnh ẩn 2D cuối cùng, bạn cần giữ latent_dim thành 2. Ngoài ra, thời gian đào tạo sẽ tăng lên khi kích thước mạng tăng lên.

Bạn cũng có thể thử triển khai VAE bằng cách sử dụng một tập dữ liệu khác, chẳng hạn như CIFAR-10.

VAE có thể được thực hiện theo nhiều kiểu khác nhau và có độ phức tạp khác nhau. Bạn có thể tìm thấy các triển khai bổ sung trong các nguồn sau:

Nếu bạn muốn tìm hiểu thêm về chi tiết của VAE, vui lòng tham khảo Giới thiệu về Bộ mã tự động biến thể .