Generative Adversarial Network

ProbFlow isn’t really built for the kind of flexibility you need to fit a GAN, so I wouldn’t recommend fitting GANs with ProbFlow. That said, you can technically do it.

So… Just for fun…

TODO: description… (Bayesian GAN, cite https://arxiv.org/pdf/1705.09558.pdf)

TODO: math

TODO: diagram

TODO: talk about overriding the Model.log_likelihood method to compute the log likelihood due to the batch of acual data and a batch of the same size of data generated by the generator network.

First let’s build a generator:

TensorFlowPyTorch

import probflow as pf
import tensorflow as tf

class Generator(pf.Model):

    def __init__(self, dims):
        self.Dz = dims[0]
        self.G = pf.DenseNetwork(dims)
        self.D = None

    def __call__(self, x):
        z = tf.random.normal([x.shape[0], self.Dz])
        return self.G(z)

    def log_likelihood(self, _, x):
        labels = tf.ones([x.shape[0], 1])
        true_ll = self.D(self(x)).log_prob(labels)
        return tf.reduce_sum(true_ll)

Then a discriminator:

TensorFlowPyTorch

class Discriminator(pf.Model):

    def __init__(self, dims):
        self.G = None
        self.D = pf.DenseNetwork(dims)

    def __call__(self, x):
        return pf.Bernoulli(self.D(x))

    def log_likelihood(self, _, x):
        labels = tf.ones([x.shape[0], 1])
        true_ll = self(x).log_prob(labels)
        fake_ll = self(self.G(x)).log_prob(0*labels)
        return tf.reduce_sum(true_ll + fake_ll)

And a callback to train the generator for an epoch at the end of each epoch of the discriminator’s training:

class TrainGenerator(pf.Callback):

    def __init__(self, G, x):
        self.G = G
        self.x = x

    def on_epoch_end(self):
        self.G.fit(self.x, epochs=1)

Then, we can instantiate the networks, the callback, and fit the network:

# x is a numpy array or pandas DataFrame of real data
Nf = 7 #number of features / input dimensionality (x.shape[1])
Nz = 3 #number of latent dimensions

# Create the networks
G = Generator([Nz, 256, 128, Nf])
D = Discriminator([Nf, 256, 128, 1])

# Let them know about each other <3
G.D = lambda x: D(x)
D.G = lambda x: G(x)

# Create the callback which trains the generator
train_g = TrainGenerator(G, x)

# Fit both models by fitting the discriminator w/ the callback
D.fit(x, callbacks=[train_g])

Note that we use lambda functions instead of simply assigning the opposite net as an attribute of the each model instance. This is because ProbFlow recursively searches a model instance’s attributes for Module and Parameter objects, and optimizes all found parameters with respect to the loss. However, we don’t want the discriminator’s parameters updated with the the generator’s loss, or vice versa! ProbFlow even looks for parameters in attributes which are lists and dictionaries, but it doesn’t look in lambda functions, so we can “hide” the parameters of one model from the other that way, while still allowing each model to __call__ the other to compute its own loss.