如何保存可学习的张量流分布？答案

【问题标题】：How to save a learnable tensorflow distribution?如何保存可学习的张量流分布？
【发布时间】：2021-01-21 10:42:00
【问题描述】：

我一直在尝试保存和加载Tensorflow distribution。似乎没有任何关于如何做到这一点的好例子。

具体来说，我一直在尝试保存继承自 tfp.distributions.Distribution 的 PixelCNN++ distribution。可以在下面找到一个最小的工作示例（取自 PixelCNN 文档）。

# Build a small Pixel CNN++ model to train on MNIST.

import tensorflow as tf
import tensorflow_datasets as tfds
import tensorflow_probability as tfp

tfd = tfp.distributions
tfk = tf.keras
tfkl = tf.keras.layers

tf.enable_v2_behavior()

# Load MNIST from tensorflow_datasets
data = tfds.load('mnist')
train_data, test_data = data['train'], data['test']

def image_preprocess(x):
  x['image'] = tf.cast(x['image'], tf.float32)
  return (x['image'],)  # (input, output) of the model

batch_size = 16
train_it = train_data.map(image_preprocess).batch(batch_size).shuffle(1000)

image_shape = (28, 28, 1)
# Define a Pixel CNN network
dist = tfd.PixelCNN(
    image_shape=image_shape,
    num_resnet=1,
    num_hierarchies=2,
    num_filters=32,
    num_logistic_mix=5,
    dropout_p=.3,
)

# Define the model input
image_input = tfkl.Input(shape=image_shape)

# Define the log likelihood for the loss fn
log_prob = dist.log_prob(image_input)

# Define the model
model = tfk.Model(inputs=image_input, outputs=log_prob)
model.add_loss(-tf.reduce_mean(log_prob))

# Compile and train the model
model.compile(
    optimizer=tfk.optimizers.Adam(.001),
    metrics=[])

model.fit(train_it, epochs=10, verbose=True)

# sample five images from the trained model
samples = dist.sample(5)

如何保存分发？或者，如何从模型中提取分布（可以是 easily saved）以便从中采样？

【问题讨论】：

标签： python tensorflow machine-learning keras

【解决方案1】：

有趣的问题，因为我目前也在尝试做类似的事情。

似乎可以以检查点样式保存权重，因为我将其用作训练过程的一部分：

import os
import time
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow.keras.callbacks import Callback

class CheckpointsCallback(Callback):
    def __init__(self, checkpoints_path):
        self.checkpoints_path = checkpoints_path
        os.makedirs(self.checkpoints_path, exist_ok=True)

    def on_epoch_end(self, epoch, logs=None):
        if self.checkpoints_path is not None:
            checkpoint_dir_path = os.path.join(self.checkpoints_path, f"epoch_{epoch}")
            checkpoint_file_path = os.path.join(checkpoint_dir_path, f"model")
            os.makedirs(checkpoint_dir_path, exist_ok=True)
            self.model.save_weights(checkpoint_file_path)
            print(f"saved weights to {checkpoint_file_path}")

callbacks = [CheckpointsCallback(checkpoints_path=checkpoints_path)]


dist.fit(x=train_dataset,
              validation_data=valid_dataset,
              epochs=epochs,
              verbose=True,
              callbacks=callbacks)

但我尝试过的 keras.Model.load_weights 或 tf.train.Checkpoint.restore 解决方案似乎都不起作用

【讨论】：