如何在tensorflow 2.0 w/keras中保存/恢复大型模型？

Question

我有一个大型定制模型与新的tensorflow 2.0和混合keras和tensorflow。我想保存它（架构和权重）。 复制的确切命令：

import tensorflow as tf


OUTPUT_CHANNELS = 3

def downsample(filters, size, apply_batchnorm=True):
  initializer = tf.random_normal_initializer(0., 0.02)

  result = tf.keras.Sequential()
  result.add(
      tf.keras.layers.Conv2D(filters, size, strides=2, padding='same',
                             kernel_initializer=initializer, use_bias=False))

  if apply_batchnorm:
    result.add(tf.keras.layers.BatchNormalization())

  result.add(tf.keras.layers.LeakyReLU())

  return result

def upsample(filters, size, apply_dropout=False):
  initializer = tf.random_normal_initializer(0., 0.02)

  result = tf.keras.Sequential()
  result.add(
    tf.keras.layers.Conv2DTranspose(filters, size, strides=2,
                                    padding='same',
                                    kernel_initializer=initializer,
                                    use_bias=False))

  result.add(tf.keras.layers.BatchNormalization())

  if apply_dropout:
      result.add(tf.keras.layers.Dropout(0.5))

  result.add(tf.keras.layers.ReLU())

  return result


def Generator():
  down_stack = [
    downsample(64, 4, apply_batchnorm=False), # (bs, 128, 128, 64)
    downsample(128, 4), # (bs, 64, 64, 128)
    downsample(256, 4), # (bs, 32, 32, 256)
    downsample(512, 4), # (bs, 16, 16, 512)
    downsample(512, 4), # (bs, 8, 8, 512)
    downsample(512, 4), # (bs, 4, 4, 512)
    downsample(512, 4), # (bs, 2, 2, 512)
    downsample(512, 4), # (bs, 1, 1, 512)
  ]

  up_stack = [
    upsample(512, 4, apply_dropout=True), # (bs, 2, 2, 1024)
    upsample(512, 4, apply_dropout=True), # (bs, 4, 4, 1024)
    upsample(512, 4, apply_dropout=True), # (bs, 8, 8, 1024)
    upsample(512, 4), # (bs, 16, 16, 1024)
    upsample(256, 4), # (bs, 32, 32, 512)
    upsample(128, 4), # (bs, 64, 64, 256)
    upsample(64, 4), # (bs, 128, 128, 128)
  ]

  initializer = tf.random_normal_initializer(0., 0.02)
  last = tf.keras.layers.Conv2DTranspose(OUTPUT_CHANNELS, 4,
                                         strides=2,
                                         padding='same',
                                         kernel_initializer=initializer,
                                         activation='tanh') # (bs, 256, 256, 3)

  concat = tf.keras.layers.Concatenate()

  inputs = tf.keras.layers.Input(shape=[None,None,3])
  x = inputs

  # Downsampling through the model
  skips = []
  for down in down_stack:
    x = down(x)
    skips.append(x)

  skips = reversed(skips[:-1])

  # Upsampling and establishing the skip connections
  for up, skip in zip(up_stack, skips):
    x = up(x)
    x = concat([x, skip])

  x = last(x)

  return tf.keras.Model(inputs=inputs, outputs=x)

generator = Generator()
generator.summary()

generator.save('generator.h5')
generator_loaded = tf.keras.models.load_model('generator.h5')

我设法用以下方式保存模型：

^{pr2}$

但当我试着加载：

generator_loaded = tf.keras.models.load_model('generator.h5')

它永远不会结束（没有错误消息）。也许模型太大了？我试图用model.to_json()和完整的API tf.keras.models.save_model()保存为JSON，但同样的问题是，无法加载它（或者至少太长了）。在

在Windows/Linux和有/没有GPU的情况下也有同样的问题。在

保存和恢复在完整的Keras和简单的模型下很好地工作。在

编辑

• 节省权重然后加载它们效果很好，但是不可能加载模型结构。在
• 我把我用来重现bug的模型放在这里，它来自Pix2Pix示例（https://www.tensorflow.org/alpha/tutorials/generative/pix2pix）
• 我还写了一篇关于tensorflow github的文章：https://github.com/tensorflow/tensorflow/issues/28281

Answer 1

我找到了一个临时的解决办法。 似乎这个问题发生在顺序API tf.keras.Sequential，通过使用函数API，tf.keras.models.load_model设法加载保存的模型。 我希望他们能在最终版本中解决这个问题，看看我在githubhttps://github.com/tensorflow/tensorflow/issues/28281中提出的问题。在

干杯

Answer 2

另一种保存经过训练的模型的方法是使用python中的pickle模块。在

import pickle
pickle.dump(model, open(filename, 'wb'))

为了加载pickled模型

^{pr2}$

pickle文件的扩展名通常是.sav

Answer 3

请尝试将模型另存为：

model.save('model_name.model')

然后加载：

^{pr2}$