How to save a learnable tensorflow distribution?

Question

I have been trying to save and load a Tensorflow distribution. There don't seem to be any good examples out there on how to do this.

Specifically I have been trying to save a PixelCNN++ distribution which inherits from tfp.distributions.Distribution. A minimum working example can be found below (this is taken from the PixelCNN documentation).

# 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)

How do I save the distribution? Or instead, how do I extract the distribution from the model (which can be easily saved) in order to sample from it?

user_2697831 · Answer

Interesting question, as I am also currently trying to do something similar.

It appears that weights can be saved, in a checkpoint style, as I use it as part of my training process:

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)

But none of the keras.Model.load_weights or tf.train.Checkpoint.restore solution that I tried seemed to work

How to save a learnable tensorflow distribution?

1 Answers