使用 CelebA 渐进式 GAN 模型生成人工面部

View 在 TensorFlow.org 上查看 在 Google Colab 中运行 在 GitHub 上查看源代码 下载笔记本 查看 TF Hub 模型

本 Colab 演示了如何使用基于生成对抗网络 (GAN) 的 TF-Hub 模块。该模块从 N 维向量(称为隐空间)映射到 RGB 图像。

本文提供了两个示例:

  • 从隐空间映射到图像,以及
  • 提供一个目标图像,利用梯度下降法找到生成与目标图像相似的图像的隐向量。

可选前提条件

更多模型

您可以在这里找到 tfhub.dev 上当前托管的所有模型,您可以使用这些模型生成图像。

设置

# Install imageio for creating animations.
pip -q install imageio
pip -q install scikit-image
pip install git+https://github.com/tensorflow/docs

Imports and function definitions

2023-11-07 18:59:27.352652: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2023-11-07 18:59:27.352699: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2023-11-07 18:59:27.354507: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered

隐空间插值法

随机向量

两个随机初始化向量之间的隐空间插值。我们将使用包含预训练渐进式 GAN 的 TF-Hub 模块 progan-128

progan = hub.load("https://tfhub.dev/google/progan-128/1").signatures['default']
def interpolate_between_vectors():
  v1 = tf.random.normal([latent_dim])
  v2 = tf.random.normal([latent_dim])

  # Creates a tensor with 25 steps of interpolation between v1 and v2.
  vectors = interpolate_hypersphere(v1, v2, 50)

  # Uses module to generate images from the latent space.
  interpolated_images = progan(vectors)['default']

  return interpolated_images

interpolated_images = interpolate_between_vectors()
animate(interpolated_images)

gif

查找隐空间中的最近向量

确定目标图像。例如,使用从模块生成的图像或上传自己的图像。

image_from_module_space = True  # @param { isTemplate:true, type:"boolean" }

def get_module_space_image():
  vector = tf.random.normal([1, latent_dim])
  images = progan(vector)['default'][0]
  return images

def upload_image():
  uploaded = files.upload()
  image = imageio.imread(uploaded[list(uploaded.keys())[0]])
  return transform.resize(image, [128, 128])

if image_from_module_space:
  target_image = get_module_space_image()
else:
  target_image = upload_image()

display_image(target_image)

png

定义目标图像与隐空间变量生成的图像之后,我们可以利用梯度下降法找到最大限度减少损失的变量值。

tf.random.set_seed(42)
initial_vector = tf.random.normal([1, latent_dim])
display_image(progan(initial_vector)['default'][0])

png

def find_closest_latent_vector(initial_vector, num_optimization_steps,
                               steps_per_image):
  images = []
  losses = []

  vector = tf.Variable(initial_vector)  
  optimizer = tf.optimizers.Adam(learning_rate=0.01)
  loss_fn = tf.losses.MeanAbsoluteError(reduction="sum")

  for step in range(num_optimization_steps):
    if (step % 100)==0:
      print()
    print('.', end='')
    with tf.GradientTape() as tape:
      image = progan(vector.read_value())['default'][0]
      if (step % steps_per_image) == 0:
        images.append(image.numpy())
      target_image_difference = loss_fn(image, target_image[:,:,:3])
      # The latent vectors were sampled from a normal distribution. We can get
      # more realistic images if we regularize the length of the latent vector to 
      # the average length of vector from this distribution.
      regularizer = tf.abs(tf.norm(vector) - np.sqrt(latent_dim))

      loss = target_image_difference + regularizer
      losses.append(loss.numpy())
    grads = tape.gradient(loss, [vector])
    optimizer.apply_gradients(zip(grads, [vector]))

  return images, losses


num_optimization_steps=200
steps_per_image=5
images, loss = find_closest_latent_vector(initial_vector, num_optimization_steps, steps_per_image)
.
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1699383604.217769   80180 device_compiler.h:186] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.
...................................................................................................
....................................................................................................
plt.plot(loss)
plt.ylim([0,max(plt.ylim())])
(0.0, 6696.204772949219)

png

animate(np.stack(images))

gif

将结果与目标进行对比:

display_image(np.concatenate([images[-1], target_image], axis=1))

png

试运行上述示例

如果图像来自模块空间,则下降很快且会收敛到合理的样本。如果尝试下降到不是来自模块空间的图像,则只有当图像相当接近训练图像的空间时,下降才会收敛。

如何使其更快速地下降并变成更真实的图像?您可以尝试:

  • 对图像差异使用不同的损失,例如二次方程,
  • 对隐向量使用不同的正则化器,
  • 在多次运行中从随机向量初始化,
  • 等等。