pot.py

# Copyright 2017 Max Planck Society
# Distributed under the BSD-3 Software license,
# (See accompanying file ./LICENSE.txt or copy at
# https://opensource.org/licenses/BSD-3-Clause)
"""This class implements POT training.

"""
import collections
import logging
import os
import time
import tensorflow as tf
import utils
from utils import ProgressBar
from utils import TQDM
import numpy as np
import ops
from metrics import Metrics
slim = tf.contrib.slim


def vgg_16(inputs,
           is_training=False,
           dropout_keep_prob=0.5,
           scope='vgg_16',
           fc_conv_padding='VALID', reuse=None):
    inputs = inputs * 255.0
    inputs -= tf.constant([123.68, 116.779, 103.939], dtype=tf.float32)
    with tf.variable_scope(scope, 'vgg_16', [inputs], reuse=reuse) as sc:
      end_points_collection = sc.name + '_end_points'
      end_points = {}
      # Collect outputs for conv2d, fully_connected and max_pool2d.
      with slim.arg_scope([slim.conv2d, slim.fully_connected, slim.max_pool2d],
                          outputs_collections=end_points_collection):
        end_points['pool0'] = inputs
        net = slim.repeat(inputs, 2, slim.conv2d, 64, [3, 3], scope='conv1')
        net = slim.max_pool2d(net, [2, 2], scope='pool1')
        end_points['pool1'] = net
        net = slim.repeat(net, 2, slim.conv2d, 128, [3, 3], scope='conv2')
        net = slim.max_pool2d(net, [2, 2], scope='pool2')
        end_points['pool2'] = net
        net = slim.repeat(net, 3, slim.conv2d, 256, [3, 3], scope='conv3')
        net = slim.max_pool2d(net, [2, 2], scope='pool3')
        end_points['pool3'] = net
        net = slim.repeat(net, 3, slim.conv2d, 512, [3, 3], scope='conv4')
        net = slim.max_pool2d(net, [2, 2], scope='pool4')
        end_points['pool4'] = net
        net = slim.repeat(net, 3, slim.conv2d, 512, [3, 3], scope='conv5')
        net = slim.max_pool2d(net, [2, 2], scope='pool5')
        end_points['pool5'] = net
  #       # Use conv2d instead of fully_connected layers.
  #       net = slim.conv2d(net, 4096, [7, 7], padding=fc_conv_padding, scope='fc6')
  #       net = slim.dropout(net, dropout_keep_prob, is_training=is_training,
  #                          scope='dropout6')
  #       net = slim.conv2d(net, 4096, [1, 1], scope='fc7')
  #       net = slim.dropout(net, dropout_keep_prob, is_training=is_training,
  #                          scope='dropout7')
  #       net = slim.conv2d(net, num_classes, [1, 1],
  #                         activation_fn=None,
  #                         normalizer_fn=None,
  #                         scope='fc8')
        # Convert end_points_collection into a end_point dict.
  #       end_points = slim.utils.convert_collection_to_dict(end_points_collection)
        return net, end_points

def compute_moments(_inputs, moments=[2, 3]):
    """From an image input, compute moments"""
    _inputs_sq = tf.square(_inputs)
    _inputs_cube = tf.pow(_inputs, 3)
    height = int(_inputs.get_shape()[1])
    width = int(_inputs.get_shape()[2])
    channels = int(_inputs.get_shape()[3])
    def ConvFlatten(x, kernel_size):
#                 w_sum = tf.ones([kernel_size, kernel_size, channels, 1]) / (kernel_size * kernel_size * channels)
        w_sum = tf.eye(num_rows=channels, num_columns=channels, batch_shape=[kernel_size * kernel_size])
        w_sum = tf.reshape(w_sum, [kernel_size, kernel_size, channels, channels])
        w_sum = w_sum / (kernel_size * kernel_size)
        sum_ = tf.nn.conv2d(x, w_sum, strides=[1, 1, 1, 1], padding='VALID')
        size = prod_dim(sum_)
        assert size == (height - kernel_size + 1) * (width - kernel_size + 1) * channels, size
        return tf.reshape(sum_, [-1, size])
    outputs = []
    for size in [3, 4, 5]:
        mean = ConvFlatten(_inputs, size)
        square = ConvFlatten(_inputs_sq, size)
        var = square - tf.square(mean)
        if 2 in moments:
            outputs.append(var)
        if 3 in moments:
            cube = ConvFlatten(_inputs_cube, size)
            skewness = cube - 3.0 * mean * var - tf.pow(mean, 3)  # Unnormalized
            outputs.append(skewness)
    return tf.concat(outputs, 1)

def prod_dim(tensor):
    return np.prod([int(d) for d in tensor.get_shape()[1:]])

def flatten(tensor):
    return tf.reshape(tensor, [-1, prod_dim(tensor)])

class Pot(object):
    """A base class for running individual POTs.

    """
    def __init__(self, opts, data, weights):

        # Create a new session with session.graph = default graph
        self._session = tf.Session()
        self._trained = False
        self._data = data
        self._data_weights = np.copy(weights)
        # Latent noise sampled ones to apply decoder while training
        self._noise_for_plots = opts['pot_pz_std'] * utils.generate_noise(opts, 1000)
        # Placeholders
        self._real_points_ph = None
        self._noise_ph = None
        # Init ops
        self._additional_init_ops = []
        self._init_feed_dict = {}

        # Main operations

        # Optimizers

        with self._session.as_default(), self._session.graph.as_default():
            logging.error('Building the graph...')
            self._build_model_internal(opts)

        # Make sure AdamOptimizer, if used in the Graph, is defined before
        # calling global_variables_initializer().
        init = tf.global_variables_initializer()
        self._session.run(init)
        self._session.run(self._additional_init_ops, self._init_feed_dict)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        # Cleaning the whole default Graph
        logging.error('Cleaning the graph...')
        tf.reset_default_graph()
        logging.error('Closing the session...')
        # Finishing the session
        self._session.close()

    def train(self, opts):
        """Train a POT model.

        """
        with self._session.as_default(), self._session.graph.as_default():
            self._train_internal(opts)
            self._trained = True

    def sample(self, opts, num=100):
        """Sample points from the trained POT model.

        """
        assert self._trained, 'Can not sample from the un-trained POT'
        with self._session.as_default(), self._session.graph.as_default():
            return self._sample_internal(opts, num)

    def train_mixture_discriminator(self, opts, fake_images):
        """Train classifier separating true data from points in fake_images.

        Return:
            prob_real: probabilities of the points from training data being the
                real points according to the trained mixture classifier.
                Numpy vector of shape (self._data.num_points,)
            prob_fake: probabilities of the points from fake_images being the
                real points according to the trained mixture classifier.
                Numpy vector of shape (len(fake_images),)

        """
        with self._session.as_default(), self._session.graph.as_default():
            return self._train_mixture_discriminator_internal(opts, fake_images)


    def _run_batch(self, opts, operation, placeholder, feed,
                   placeholder2=None, feed2=None):
        """Wrapper around session.run to process huge data.

        It is asumed that (a) first dimension of placeholder enumerates
        separate points, and (b) that operation is independently applied
        to every point, i.e. we can split it point-wisely and then merge
        the results. The second placeholder is meant either for is_train
        flag for batch-norm or probabilities of dropout.

        TODO: write util function which will be called both from this method
        and MNIST classification evaluation as well.

        """
        assert len(feed.shape) > 0, 'Empry feed.'
        num_points = feed.shape[0]
        batch_size = opts['tf_run_batch_size']
        batches_num = int(np.ceil((num_points + 0.) / batch_size))
        result = []
        for idx in xrange(batches_num):
            if idx == batches_num - 1:
                if feed2 is None:
                    res = self._session.run(
                        operation,
                        feed_dict={placeholder: feed[idx * batch_size:]})
                else:
                    res = self._session.run(
                        operation,
                        feed_dict={placeholder: feed[idx * batch_size:],
                                   placeholder2: feed2})
            else:
                if feed2 is None:
                    res = self._session.run(
                        operation,
                        feed_dict={placeholder: feed[idx * batch_size:
                                                     (idx + 1) * batch_size]})
                else:
                    res = self._session.run(
                        operation,
                        feed_dict={placeholder: feed[idx * batch_size:
                                                     (idx + 1) * batch_size],
                                   placeholder2: feed2})

            if len(res.shape) == 1:
                # convert (n,) vector to (n,1) array
                res = np.reshape(res, [-1, 1])
            result.append(res)
        result = np.vstack(result)
        assert len(result) == num_points
        return result

    def _build_model_internal(self, opts):
        """Build a TensorFlow graph with all the necessary ops.

        """
        assert False, 'POT base class has no build_model method defined.'

    def _train_internal(self, opts):
        assert False, 'POT base class has no train method defined.'

    def _sample_internal(self, opts, num):
        assert False, 'POT base class has no sample method defined.'

    def _train_mixture_discriminator_internal(self, opts, fake_images):
        assert False, 'POT base class has no mixture discriminator method defined.'


class ImagePot(Pot):
    """A simple POT implementation, suitable for pictures.

    """

    def __init__(self, opts, data, weights):

        # One more placeholder for batch norm
        self._is_training_ph = None
        Pot.__init__(self, opts, data, weights)


    def dcgan_like_arch(self, opts, noise, is_training, reuse, keep_prob):
        output_shape = self._data.data_shape
        num_units = opts['g_num_filters']

        batch_size = tf.shape(noise)[0]
        num_layers = opts['g_num_layers']
        if opts['g_arch'] == 'dcgan':
            height = output_shape[0] / 2**num_layers
            width = output_shape[1] / 2**num_layers
        elif opts['g_arch'] == 'dcgan_mod':
            height = output_shape[0] / 2**(num_layers-1)
            width = output_shape[1] / 2**(num_layers-1)
        else:
            assert False

        h0 = ops.linear(
            opts, noise, num_units * height * width, scope='h0_lin')
        h0 = tf.reshape(h0, [-1, height, width, num_units])
        h0 = tf.nn.relu(h0)
        layer_x = h0
        for i in xrange(num_layers-1):
            scale = 2**(i+1)
            if opts['g_stride1_deconv']:
                # Sylvain, I'm worried about this part!
                _out_shape = [batch_size, height * scale / 2,
                              width * scale / 2, num_units / scale * 2]
                layer_x = ops.deconv2d(
                    opts, layer_x, _out_shape, d_h=1, d_w=1,
                    scope='h%d_deconv_1x1' % i)
                layer_x = tf.nn.relu(layer_x)
            _out_shape = [batch_size, height * scale, width * scale, num_units / scale]
            layer_x = ops.deconv2d(opts, layer_x, _out_shape, scope='h%d_deconv' % i)
            if opts['batch_norm']:
                layer_x = ops.batch_norm(opts, layer_x, is_training, reuse, scope='bn%d' % i)
            layer_x = tf.nn.relu(layer_x)
            if opts['dropout']:
                _keep_prob = tf.minimum(
                    1., 0.9 - (0.9 - keep_prob) * float(i + 1) / (num_layers - 1))
                layer_x = tf.nn.dropout(layer_x, _keep_prob)

        _out_shape = [batch_size] + list(output_shape)
        if opts['g_arch'] == 'dcgan':
            last_h = ops.deconv2d(
                opts, layer_x, _out_shape, scope='hlast_deconv')
        elif opts['g_arch'] == 'dcgan_mod':
            last_h = ops.deconv2d(
                opts, layer_x, _out_shape, d_h=1, d_w=1, scope='hlast_deconv')
        else:
            assert False

        if opts['input_normalize_sym']:
            return tf.nn.tanh(last_h)
        else:
            return tf.nn.sigmoid(last_h)

    def began_dec(self, opts, noise, is_training, reuse, keep_prob):
        """ Architecture reported here: https://arxiv.org/pdf/1703.10717.pdf
        """

        output_shape = self._data.data_shape
        num_units = opts['g_num_filters']
        num_layers = opts['g_num_layers']
        batch_size = tf.shape(noise)[0]

        h0 = ops.linear(
            opts, noise, num_units * 8 * 8, scope='h0_lin')
        h0 = tf.reshape(h0, [-1, 8, 8, num_units])
        layer_x = h0
        for i in xrange(num_layers):
            if i % 3 < 2:
                # Don't change resolution
                layer_x = ops.conv2d(opts, layer_x, num_units, d_h=1, d_w=1, scope='h%d_conv' % i)
                layer_x = tf.nn.elu(layer_x)
            else:
                if i != num_layers - 1:
                    # Upsampling by factor of 2 with NN
                    scale = 2 ** (i / 3 + 1)
                    layer_x = ops.upsample_nn(layer_x, [scale * 8, scale * 8],
                                              scope='h%d_upsample' % i, reuse=reuse)
                    # Skip connection
                    append = ops.upsample_nn(h0, [scale * 8, scale * 8],
                                              scope='h%d_skipup' % i, reuse=reuse)
                    layer_x = tf.concat([layer_x, append], axis=3)

        last_h = ops.conv2d(opts, layer_x, output_shape[-1], d_h=1, d_w=1, scope='hlast_conv')

        if opts['input_normalize_sym']:
            return tf.nn.tanh(last_h)
        else:
            return tf.nn.sigmoid(last_h)

    def conv_up_res(self, opts, noise, is_training, reuse, keep_prob):
        output_shape = self._data.data_shape
        num_units = opts['g_num_filters']

        batch_size = tf.shape(noise)[0]
        num_layers = opts['g_num_layers']
        data_height = output_shape[0]
        data_width = output_shape[1]
        data_channels = output_shape[2]
        height = data_height / 2**num_layers
        width = data_width / 2**num_layers

        h0 = ops.linear(
            opts, noise, num_units * height * width, scope='h0_lin')
        h0 = tf.reshape(h0, [-1, height, width, num_units])
        h0 = tf.nn.relu(h0)
        layer_x = h0
        for i in xrange(num_layers-1):
            layer_x = tf.image.resize_nearest_neighbor(layer_x, (2 * height, 2 * width))
            layer_x = ops.conv2d(opts, layer_x, num_units / 2, d_h=1, d_w=1, scope='conv2d_%d' % i)
            height *= 2
            width *= 2
            num_units /= 2

            if opts['g_3x3_conv'] > 0:
                before = layer_x
                for j in range(opts['g_3x3_conv']):
                    layer_x = ops.conv2d(opts, layer_x, num_units, d_h=1, d_w=1,
                                         scope='conv2d_3x3_%d_%d' % (i, j),
                                         conv_filters_dim=3)
                    layer_x = tf.nn.relu(layer_x)
                layer_x += before  # Residual connection.

            if opts['batch_norm']:
                layer_x = ops.batch_norm(opts, layer_x, is_training, reuse, scope='bn%d' % i)
            layer_x = tf.nn.relu(layer_x)
            if opts['dropout']:
                _keep_prob = tf.minimum(
                    1., 0.9 - (0.9 - keep_prob) * float(i + 1) / (num_layers - 1))
                layer_x = tf.nn.dropout(layer_x, _keep_prob)

        layer_x = tf.image.resize_nearest_neighbor(layer_x, (2 * height, 2 * width))
        layer_x = ops.conv2d(opts, layer_x, data_channels, d_h=1, d_w=1, scope='last_conv2d_%d' % i)

        if opts['input_normalize_sym']:
            return tf.nn.tanh(layer_x)
        else:
            return tf.nn.sigmoid(layer_x)

    def ali_deconv(self, opts, noise, is_training, reuse, keep_prob):
        output_shape = self._data.data_shape

        batch_size = tf.shape(noise)[0]
        noise_size = int(noise.get_shape()[1])
        data_height = output_shape[0]
        data_width = output_shape[1]
        data_channels = output_shape[2]

        noise = tf.reshape(noise, [-1, 1, 1, noise_size])

        num_units = opts['g_num_filters']
        layer_params = []
        layer_params.append([4, 1, num_units])
        layer_params.append([4, 2, num_units / 2])
        layer_params.append([4, 1, num_units / 4])
        layer_params.append([4, 2, num_units / 8])
        layer_params.append([5, 1, num_units / 8])
        # For convolution: (n - k) / stride + 1 = s
        # For transposed: (s - 1) * stride + k = n
        layer_x = noise
        height = 1
        width = 1
        for i, (kernel, stride, channels) in enumerate(layer_params):
            height = (height - 1) * stride + kernel
            width = height
            layer_x = ops.deconv2d(
                opts, layer_x, [batch_size, height, width, channels], d_h=stride, d_w=stride,
                scope='h%d_deconv' % i, conv_filters_dim=kernel, padding='VALID')
            if opts['batch_norm']:
                layer_x = ops.batch_norm(opts, layer_x, is_training, reuse, scope='bn%d' % i)
            layer_x = ops.lrelu(layer_x, 0.1)
        assert height == data_height
        assert width == data_width

        # Then two 1x1 convolutions.
        layer_x = ops.conv2d(opts, layer_x, num_units / 8, d_h=1, d_w=1, scope='conv2d_1x1', conv_filters_dim=1)
        if opts['batch_norm']:
            layer_x = ops.batch_norm(opts, layer_x, is_training, reuse, scope='bnlast')
        layer_x = ops.lrelu(layer_x, 0.1)
        layer_x = ops.conv2d(opts, layer_x, data_channels, d_h=1, d_w=1, scope='conv2d_1x1_2', conv_filters_dim=1)

        if opts['input_normalize_sym']:
            return tf.nn.tanh(layer_x)
        else:
            return tf.nn.sigmoid(layer_x)

    def generator(self, opts, noise, is_training=False, reuse=False, keep_prob=1.):
        """ Decoder actually.

        """

        output_shape = self._data.data_shape
        num_units = opts['g_num_filters']

        with tf.variable_scope("GENERATOR", reuse=reuse):
            # if not opts['convolutions']:
            if opts['g_arch'] == 'mlp':
                layer_x = noise
                for i in range(opts['g_num_layers']):
                    layer_x = ops.linear(opts, layer_x, num_units, 'h%d_lin' % i)
                    layer_x = tf.nn.relu(layer_x)
                    if opts['batch_norm']:
                        layer_x = ops.batch_norm(
                            opts, layer_x, is_training, reuse, scope='bn%d' % i)
                out = ops.linear(opts, layer_x, np.prod(output_shape), 'h%d_lin' % (i + 1))
                out = tf.reshape(out, [-1] + list(output_shape))
                if opts['input_normalize_sym']:
                    return tf.nn.tanh(out)
                else:
                    return tf.nn.sigmoid(out)
            elif opts['g_arch'] in ['dcgan', 'dcgan_mod']:
                return self.dcgan_like_arch(opts, noise, is_training, reuse, keep_prob)
            elif opts['g_arch'] == 'conv_up_res':
                return self.conv_up_res(opts, noise, is_training, reuse, keep_prob)
            elif opts['g_arch'] == 'ali':
                return self.ali_deconv(opts, noise, is_training, reuse, keep_prob)
            elif opts['g_arch'] == 'began':
                return self.began_dec(opts, noise, is_training, reuse, keep_prob)
            else:
                raise ValueError('%s unknown' % opts['g_arch'])

    def discriminator(self, opts, input_, prefix='DISCRIMINATOR', reuse=False):
        """Discriminator for the GAN objective

        """
        num_units = opts['d_num_filters']
        num_layers = opts['d_num_layers']
        nowozin_trick = opts['gan_p_trick']
        # No convolutions as GAN happens in the latent space
        with tf.variable_scope(prefix, reuse=reuse):
            hi = input_
            for i in range(num_layers):
                hi = ops.linear(opts, hi, num_units, scope='h%d_lin' % (i+1))
                hi = tf.nn.relu(hi)
            hi = ops.linear(opts, hi, 1, scope='final_lin')
        if nowozin_trick:
            # We are doing GAN between our model Qz and the true Pz.
            # We know analytical form of the true Pz.
            # The optimal discriminator for D_JS(Pz, Qz) is given by:
            # Dopt(x) = log dPz(x) - log dQz(x)
            # And we know exactly dPz(x). So add log dPz(x) explicitly 
            # to the discriminator and let it learn only the remaining
            # dQz(x) term. This appeared in the AVB paper.
            assert opts['latent_space_distr'] == 'normal'
            sigma2_p = float(opts['pot_pz_std']) ** 2
            normsq = tf.reduce_sum(tf.square(input_), 1)
            hi = hi - normsq / 2. / sigma2_p \
                    - 0.5 * tf.log(2. * np.pi) \
                    - 0.5 * opts['latent_space_dim'] * np.log(sigma2_p)
        return hi

    def pz_sampler(self, opts, input_, prefix='PZ_SAMPLER', reuse=False):
        """Transformation to be applied to the sample from Pz
        We are trying to match Qz to phi(Pz), where phi is defined by
        this function
        """
        dim = opts['latent_space_dim']
        with tf.variable_scope(prefix, reuse=reuse):
            matrix = tf.get_variable(
                "W", [dim, dim], tf.float32,
                tf.constant_initializer(np.identity(dim)))
            bias = tf.get_variable(
                "b", [dim],
                initializer=tf.constant_initializer(0.))
        return tf.matmul(input_, matrix) + bias

    def get_batch_size(self, opts, input_):
        return tf.cast(tf.shape(input_)[0], tf.float32)# opts['batch_size']

    def moments_stats(self, opts, input_):
        """
        Compute estimates of the first 4 moments of the coordinates
        based on the sample in input_. Compare them to the desired
        population values and return a corresponding loss.
        """
        input_ = input_ / float(opts['pot_pz_std'])
        # If Pz = Qz then input_ should now come from 
        # a product of pz_dim Gaussians N(0, 1)
        # Thus first moments should be 0
        p1 = tf.reduce_mean(input_, 0)
        center_inp = input_ - p1 # Broadcasting
        # Second centered and normalized moments should be 1
        p2 = tf.sqrt(1e-5 + tf.reduce_mean(tf.square(center_inp), 0))
        normed_inp = center_inp / p2
        # Third central moment should be 0
        # p3 = tf.pow(1e-5 + tf.abs(tf.reduce_mean(tf.pow(center_inp, 3), 0)), 1.0 / 3.0)
        p3 = tf.abs(tf.reduce_mean(tf.pow(center_inp, 3), 0))
        # 4th central moment of any uni-variate Gaussian = 3 * sigma^4
        # p4 = tf.pow(1e-5 + tf.reduce_mean(tf.pow(center_inp, 4), 0) / 3.0, 1.0 / 4.0)
        p4 = tf.reduce_mean(tf.pow(center_inp, 4), 0) / 3.
        def zero_t(v):
            return tf.sqrt(1e-5 + tf.reduce_mean(tf.square(v)))
        def one_t(v):
            # The function below takes its minimum value 1. at v = 1.
            return tf.sqrt(1e-5 + tf.reduce_mean(tf.maximum(tf.square(v), 1.0 / (1e-5 + tf.square(v)))))
        return tf.stack([zero_t(p1), one_t(p2), zero_t(p3), one_t(p4)])

    def discriminator_test(self, opts, input_):
        """Deterministic discriminator using simple tests."""
        if opts['z_test'] == 'cramer':
            test_v = self.discriminator_cramer_test(opts, input_)
        elif opts['z_test'] == 'anderson':
            test_v = self.discriminator_anderson_test(opts, input_)
        elif opts['z_test'] == 'moments':
            test_v = tf.reduce_mean(self.moments_stats(opts, input_)) / 10.0
        elif opts['z_test'] == 'lks':
            test_v = self.discriminator_lks_test(opts, input_)
        else:
            raise ValueError('%s Unknown' % opts['z_test'])
        return test_v

    def discriminator_cramer_test(self, opts, input_):
        """Deterministic discriminator using Cramer von Mises Test.

        """
        add_dim = opts['z_test_proj_dim']
        if add_dim > 0:
            dim = int(input_.get_shape()[1])
            proj = np.random.rand(dim, add_dim)
            proj = proj - np.mean(proj, 0)
            norms = np.sqrt(np.sum(np.square(proj), 0) + 1e-5)
            proj = tf.constant(proj / norms, dtype=tf.float32)
            projected_x = tf.matmul(input_, proj)  # Shape [batch_size, add_dim].

            # Shape [batch_size, z_dim+add_dim]
            all_dims_x = tf.concat([input_, projected_x], 1)
        else:
            all_dims_x = input_

        # top_k can only sort on the last dimension and we want to sort the
        # first one (batch_size).
        batch_size = self.get_batch_size(opts, all_dims_x)
        transposed = tf.transpose(all_dims_x, perm=[1, 0])
        values, indices = tf.nn.top_k(transposed, k=tf.cast(batch_size, tf.int32))
        values = tf.reverse(values, [1])
        #values = tf.Print(values, [values], "sorted values")
        normal_dist = tf.contrib.distributions.Normal(0., float(opts['pot_pz_std']))
        #
        normal_cdf = normal_dist.cdf(values)
        #normal_cdf = tf.Print(normal_cdf, [normal_cdf], "normal_cdf")
        expected = (2 * tf.range(1, batch_size+1, 1, dtype="float") - 1) / (2.0 * batch_size)
        #expected = tf.Print(expected, [expected], "expected")
        # We don't use the constant.
        # constant = 1.0 / (12.0 * batch_size * batch_size)
        # stat = constant + tf.reduce_sum(tf.square(expected - normal_cdf), 1) / batch_size
        stat = tf.reduce_sum(tf.square(expected - normal_cdf), 1) / batch_size
        stat = tf.reduce_mean(stat)
        #stat = tf.Print(stat, [stat], "stat")
        return stat

    def discriminator_anderson_test(self, opts, input_):
        """Deterministic discriminator using the Anderson Darling test.

        """
        # A-D test says to normalize data before computing the statistic
        # Because true mean and variance are known, we are supposed to use
        # the population parameters for that, but wiki says it's better to
        # still use the sample estimates while normalizing
        means = tf.reduce_mean(input_, 0)
        input_ = input_ - means # Broadcasting
        stds = tf.sqrt(1e-5 + tf.reduce_mean(tf.square(input_), 0))
        input_= input_ / stds
        # top_k can only sort on the last dimension and we want to sort the
        # first one (batch_size).
        batch_size = self.get_batch_size(opts, input_)
        transposed = tf.transpose(input_, perm=[1, 0])
        values, indices = tf.nn.top_k(transposed, k=tf.cast(batch_size, tf.int32))
        values = tf.reverse(values, [1])
        normal_dist = tf.contrib.distributions.Normal(0., float(opts['pot_pz_std']))
        normal_cdf = normal_dist.cdf(values)
        # ln_normal_cdf is of shape (z_dim, batch_size)
        ln_normal_cdf = tf.log(normal_cdf)
        ln_one_normal_cdf = tf.log(1.0 - normal_cdf)
        w1 = 2 * tf.range(1, batch_size + 1, 1, dtype="float") - 1
        w2 = 2 * tf.range(batch_size - 1, -1, -1, dtype="float") + 1
        stat = -batch_size - tf.reduce_sum(w1 * ln_normal_cdf + \
                                           w2 * ln_one_normal_cdf, 1) / batch_size
        # stat is of shape (z_dim)
        stat = tf.reduce_mean(tf.square(stat))
        return stat

    def discriminator_lks_lin_test(self, opts, input_):
        """Deterministic discriminator using Kernel Stein Discrepancy test
        refer to LKS test on page 3 of https://arxiv.org/pdf/1705.07673.pdf

        The statistic basically reads:
            \[
                \frac{2}{n}\sum_{i=1}^n \left(
                    frac{<x_{2i}, x_{2i - 1}>}{\sigma_p^4}
                    + d/\sigma_k^2
                    - \|x_{2i} - x_{2i - 1}\|^2\left(\frac{1}{\sigma_p^2\sigma_k^2} + \frac{1}{\sigma_k^4}\right)
                \right)
                \exp( - \|x_{2i} - x_{2i - 1}\|^2/2/\sigma_k^2)
            \]

        """
        # To check the typical sizes of the test for Pz = Qz, uncomment
        # input_ = opts['pot_pz_std'] * utils.generate_noise(opts, 100000)
        batch_size = self.get_batch_size(opts, input_)
        batch_size = tf.cast(batch_size, tf.int32)
        half_size = batch_size / 2
        # s1 = tf.slice(input_, [0, 0], [half_size, -1])
        # s2 = tf.slice(input_, [half_size, 0], [half_size, -1])
        s1 = input_[:half_size, :]
        s2 = input_[half_size:, :]
        dotprods = tf.reduce_sum(tf.multiply(s1, s2), axis=1)
        distances = tf.reduce_sum(tf.square(s1 - s2), axis=1)
        sigma2_p = opts['pot_pz_std'] ** 2 # var = std ** 2
        # Median heuristic for the sigma^2 of Gaussian kernel
        # sigma2_k = tf.nn.top_k(distances, half_size).values[half_size - 1]
        # Maximum heuristic for the sigma^2 of Gaussian kernel
        # sigma2_k = tf.nn.top_k(distances, 1).values[0]
        sigma2_k = opts['latent_space_dim'] * sigma2_p
        if opts['verbose'] == 2:
            sigma2_k = tf.Print(sigma2_k, [tf.nn.top_k(distances, 1).values[0]],
                                'Maximal squared pairwise distance:')
            sigma2_k = tf.Print(sigma2_k, [tf.reduce_mean(distances)],
                                'Average squared pairwise distance:')
            sigma2_k = tf.Print(sigma2_k, [sigma2_k], 'Kernel width:')
        # sigma2_k = tf.Print(sigma2_k, [sigma2_k], 'Kernel width:')
        res = dotprods / sigma2_p ** 2 \
              - distances * (1. / sigma2_p / sigma2_k + 1. / sigma2_k ** 2) \
              + opts['latent_space_dim'] / sigma2_k
        res = tf.multiply(res, tf.exp(- distances / 2./ sigma2_k))
        stat = tf.reduce_mean(res)
        return stat


    def discriminator_lks_test(self, opts, input_):
        """Deterministic discriminator using Kernel Stein Discrepancy test
        refer to the quadratic test of https://arxiv.org/pdf/1705.07673.pdf

        The statistic basically reads:
            \[
                \frac{1}{n^2 - n}\sum_{i \neq j} \left(
                    frac{<x_i, x__j>}{\sigma_p^4}
                    + d/\sigma_k^2
                    - \|x_i - x_j\|^2\left(\frac{1}{\sigma_p^2\sigma_k^2} + \frac{1}{\sigma_k^4}\right)
                \right)
                \exp( - \|x_i - x_j\|^2/2/\sigma_k^2)
            \]

        """
        n = self.get_batch_size(opts, input_)
        n = tf.cast(n, tf.int32)
        half_size = (n * n - n) / 2
        nf = tf.cast(n, tf.float32)
        norms = tf.reduce_sum(tf.square(input_), axis=1, keep_dims=True)
        dotprods = tf.matmul(input_, input_, transpose_b=True)
        distances = norms + tf.transpose(norms) - 2. * dotprods
        sigma2_p = opts['pot_pz_std'] ** 2 # var = std ** 2
        # Median heuristic for the sigma^2 of Gaussian kernel
        # sigma2_k = tf.nn.top_k(tf.reshape(distances, [-1]), half_size).values[half_size - 1]
        # Maximal heuristic for the sigma^2 of Gaussian kernel
        # sigma2_k = tf.nn.top_k(tf.reshape(distances, [-1]), 1).values[0]
        sigma2_k = opts['latent_space_dim'] * sigma2_p
        if opts['verbose'] == 2:
            sigma2_k = tf.Print(sigma2_k, [tf.nn.top_k(tf.reshape(distances, [-1]), 1).values[0]],
                                'Maximal squared pairwise distance:')
            sigma2_k = tf.Print(sigma2_k, [tf.reduce_mean(distances)],
                                'Average squared pairwise distance:')
            sigma2_k = tf.Print(sigma2_k, [sigma2_k], 'Kernel width:')
        res = dotprods / sigma2_p ** 2 \
              - distances * (1. / sigma2_p / sigma2_k + 1. / sigma2_k ** 2) \
              + opts['latent_space_dim'] / sigma2_k
        res = tf.multiply(res, tf.exp(- distances / 2./ sigma2_k))
        res = tf.multiply(res, 1. - tf.eye(n))
        stat = tf.reduce_sum(res) / (nf * nf - nf)
        # stat = tf.reduce_sum(res) / (nf * nf)
        return stat

    def discriminator_mmd_test(self, opts, sample_qz, sample_pz):
        """U statistic for MMD(Qz, Pz) with the RBF kernel

        """
        sigma2_p = opts['pot_pz_std'] ** 2 # var = std ** 2
        kernel = 'IM'
        n = self.get_batch_size(opts, sample_qz)
        n = tf.cast(n, tf.int32)
        nf = tf.cast(n, tf.float32)
        half_size = (n * n - n) / 2
        # Pz
        norms_pz = tf.reduce_sum(tf.square(sample_pz), axis=1, keep_dims=True)
        dotprods_pz = tf.matmul(sample_pz, sample_pz, transpose_b=True)
        distances_pz = norms_pz + tf.transpose(norms_pz) - 2. * dotprods_pz
        # Qz
        norms_qz = tf.reduce_sum(tf.square(sample_qz), axis=1, keep_dims=True)
        dotprods_qz = tf.matmul(sample_qz, sample_qz, transpose_b=True)
        distances_qz = norms_qz + tf.transpose(norms_qz) - 2. * dotprods_qz
        # Pz vs Qz
        dotprods = tf.matmul(sample_qz, sample_pz, transpose_b=True)
        distances = norms_qz + tf.transpose(norms_pz) - 2. * dotprods


        if opts['verbose'] == 2:
            distances = tf.Print(distances, [tf.nn.top_k(tf.reshape(distances_qz, [-1]), 1).values[0]],
                                'Maximal Qz squared pairwise distance:')
            distances = tf.Print(distances, [tf.reduce_mean(distances_qz)],
                                'Average Qz squared pairwise distance:')

            distances = tf.Print(distances, [tf.nn.top_k(tf.reshape(distances_pz, [-1]), 1).values[0]],
                                'Maximal Pz squared pairwise distance:')
            distances = tf.Print(distances, [tf.reduce_mean(distances_pz)],
                                'Average Pz squared pairwise distance:')

            distances = tf.Print(distances, [tf.nn.top_k(tf.reshape(distances, [-1]), 1).values[0]],
                                'Maximal squared pairwise distance:')
            distances = tf.Print(distances, [tf.nn.top_k(tf.reshape(distances, [-1]), n * n).values[n * n - 1]],
                                'Minimal squared pairwise distance:')
            distances = tf.Print(distances, [tf.reduce_mean(distances)],
                                'Average squared pairwise distance:')

        if kernel == 'RBF':
            # RBF kernel

            # Median heuristic for the sigma^2 of Gaussian kernel
            sigma2_k = tf.nn.top_k(tf.reshape(distances, [-1]), half_size).values[half_size - 1]
            sigma2_k += tf.nn.top_k(tf.reshape(distances_qz, [-1]), half_size).values[half_size - 1]
            # Maximal heuristic for the sigma^2 of Gaussian kernel
            # sigma2_k = tf.nn.top_k(tf.reshape(distances_qz, [-1]), 1).values[0]
            # sigma2_k += tf.nn.top_k(tf.reshape(distances, [-1]), 1).values[0]
            # sigma2_k = opts['latent_space_dim'] * sigma2_p
            sigma2_k = tf.Print(sigma2_k, [sigma2_k], 'Kernel width:')
            res1 = tf.exp( - distances_qz / 2. / sigma2_k)
            res1 += tf.exp( - distances_pz / 2. / sigma2_k)
            res1 = tf.multiply(res1, 1. - tf.eye(n))
            res1 = tf.reduce_sum(res1) / (nf * nf - nf)
            res2 = tf.exp( - distances / 2. / sigma2_k)
            res2 = tf.reduce_sum(res2) * 2. / (nf * nf)
            stat = res1 - res2
            # stat = tf.reduce_sum(res) / (nf * nf)
        elif kernel == 'IM':
            # C = tf.nn.top_k(tf.reshape(distances, [-1]), half_size).values[half_size - 1]
            # C += tf.nn.top_k(tf.reshape(distances_qz, [-1]), half_size).values[half_size - 1]
            C = 2 * opts['latent_space_dim'] * sigma2_p
            res1 = C / (C + distances_qz)
            res1 += C / (C + distances_pz)
            res1 = tf.multiply(res1, 1. - tf.eye(n))
            res1 = tf.reduce_sum(res1) / (nf * nf - nf)
            res1 = tf.Print(res1, [res1], 'First two terms:')
            res2 = C / (C + distances)
            res2 = tf.reduce_sum(res2) * 2. / (nf * nf)
            res2 = tf.Print(res2, [res2], 'Negative term:')
            stat = res1 - res2
            # stat = tf.reduce_sum(res) / (nf * nf)
        return stat

    def correlation_loss(self, opts, input_):
        """
        Independence test based on Pearson's correlation.
        Keep in mind that this captures only linear dependancies.
        However, for multivariate Gaussian independence is equivalent
        to zero correlation.
        """

        batch_size = self.get_batch_size(opts, input_)
        dim = int(input_.get_shape()[1])
        transposed = tf.transpose(input_, perm=[1, 0])
        mean = tf.reshape(tf.reduce_mean(transposed, axis=1), [-1, 1])
        centered_transposed = transposed - mean # Broadcasting mean
        cov = tf.matmul(centered_transposed, centered_transposed, transpose_b=True)
        cov = cov / (batch_size - 1)
        #cov = tf.Print(cov, [cov], "cov")
        sigmas = tf.sqrt(tf.diag_part(cov) + 1e-5)
        #sigmas = tf.Print(sigmas, [sigmas], "sigmas")
        sigmas = tf.reshape(sigmas, [1, -1])
        sigmas = tf.matmul(sigmas, sigmas, transpose_a=True)
        #sigmas = tf.Print(sigmas, [sigmas], "sigmas")
        # Pearson's correlation
        corr = cov / sigmas
        triangle = tf.matrix_set_diag(tf.matrix_band_part(corr, 0, -1), tf.zeros(dim))
        #triangle = tf.Print(triangle, [triangle], "triangle")
        loss = tf.reduce_sum(tf.square(triangle)) / ((dim * dim - dim) / 2.0)
        #loss = tf.Print(loss, [loss], "Correlation loss")
        return loss


    def encoder(self, opts, input_, is_training=False, reuse=False, keep_prob=1.):
        if opts['e_add_noise']:
            def add_noise(x):
                shape = tf.shape(x)
                return x + tf.truncated_normal(shape, 0.0, 0.01)
            def do_nothing(x):
                return x
            input_ = tf.cond(is_training, lambda: add_noise(input_), lambda: do_nothing(input_))
        num_units = opts['e_num_filters']
        num_layers = opts['e_num_layers']
        with tf.variable_scope("ENCODER", reuse=reuse):
            if not opts['convolutions']:
                hi = input_
                for i in range(num_layers):
                    hi = ops.linear(opts, hi, num_units, scope='h%d_lin' % i)
                    if opts['batch_norm']:
                        hi = ops.batch_norm(opts, hi, is_training, reuse, scope='bn%d' % i)
                    hi = tf.nn.relu(hi)
                if opts['e_is_random']:
                    latent_mean = ops.linear(
                        opts, hi, opts['latent_space_dim'], 'h%d_lin' % (i + 1))
                    log_latent_sigmas = ops.linear(
                        opts, hi, opts['latent_space_dim'], 'h%d_lin_sigma' % (i + 1))
                    return latent_mean, log_latent_sigmas
                else:
                    return ops.linear(opts, hi, opts['latent_space_dim'], 'h%d_lin' % (i + 1))
            elif opts['e_arch'] == 'dcgan':
                return self.dcgan_encoder(opts, input_, is_training, reuse, keep_prob)
            elif opts['e_arch'] == 'ali':
                return self.ali_encoder(opts, input_, is_training, reuse, keep_prob)
            elif opts['e_arch'] == 'began':
                return self.began_encoder(opts, input_, is_training, reuse, keep_prob)
            else:
                raise ValueError('%s Unknown' % opts['e_arch'])

    def dcgan_encoder(self, opts, input_, is_training=False, reuse=False, keep_prob=1.):
        num_units = opts['e_num_filters']
        num_layers = opts['e_num_layers']
        layer_x = input_
        for i in xrange(num_layers):
            scale = 2**(num_layers-i-1)
            layer_x = ops.conv2d(opts, layer_x, num_units / scale, scope='h%d_conv' % i)

            if opts['batch_norm']:
                layer_x = ops.batch_norm(opts, layer_x, is_training, reuse, scope='bn%d' % i)
            layer_x = tf.nn.relu(layer_x)
            if opts['dropout']:
                _keep_prob = tf.minimum(
                    1., 0.9 - (0.9 - keep_prob) * float(i + 1) / num_layers)
                layer_x = tf.nn.dropout(layer_x, _keep_prob)

            if opts['e_3x3_conv'] > 0:
                before = layer_x
                for j in range(opts['e_3x3_conv']):
                    layer_x = ops.conv2d(opts, layer_x, num_units / scale, d_h=1, d_w=1,
                                         scope='conv2d_3x3_%d_%d' % (i, j),
                                         conv_filters_dim=3)
                    layer_x = tf.nn.relu(layer_x)
                layer_x += before  # Residual connection.

        if opts['e_is_random']:
            latent_mean = ops.linear(
                opts, layer_x, opts['latent_space_dim'], scope='hlast_lin')
            log_latent_sigmas = ops.linear(
                opts, layer_x, opts['latent_space_dim'], scope='hlast_lin_sigma')
            return latent_mean, log_latent_sigmas
        else:
            return ops.linear(opts, layer_x, opts['latent_space_dim'], scope='hlast_lin')

    def ali_encoder(self, opts, input_, is_training=False, reuse=False, keep_prob=1.):
        num_units = opts['e_num_filters']
        layer_params = []
        layer_params.append([5, 1, num_units / 8])
        layer_params.append([4, 2, num_units / 4])
        layer_params.append([4, 1, num_units / 2])
        layer_params.append([4, 2, num_units])
        layer_params.append([4, 1, num_units * 2])
        # For convolution: (n - k) / stride + 1 = s
        # For transposed: (s - 1) * stride + k = n
        layer_x = input_
        height = int(layer_x.get_shape()[1])
        width = int(layer_x.get_shape()[2])
        assert height == width
        for i, (kernel, stride, channels) in enumerate(layer_params):
            height = (height - kernel) / stride + 1
            width = height
            # print((height, width))
            layer_x = ops.conv2d(
                opts, layer_x, channels, d_h=stride, d_w=stride,
                scope='h%d_conv' % i, conv_filters_dim=kernel, padding='VALID')
            if opts['batch_norm']:
                layer_x = ops.batch_norm(opts, layer_x, is_training, reuse, scope='bn%d' % i)
            layer_x = ops.lrelu(layer_x, 0.1)
        assert height == 1
        assert width == 1

        # Then two 1x1 convolutions.
        layer_x = ops.conv2d(opts, layer_x, num_units * 2, d_h=1, d_w=1, scope='conv2d_1x1', conv_filters_dim=1)
        if opts['batch_norm']:
            layer_x = ops.batch_norm(opts, layer_x, is_training, reuse, scope='bnlast')
        layer_x = ops.lrelu(layer_x, 0.1)
        layer_x = ops.conv2d(opts, layer_x, num_units / 2, d_h=1, d_w=1, scope='conv2d_1x1_2', conv_filters_dim=1)

        if opts['e_is_random']:
            latent_mean = ops.linear(
                opts, layer_x, opts['latent_space_dim'], scope='hlast_lin')
            log_latent_sigmas = ops.linear(
                opts, layer_x, opts['latent_space_dim'], scope='hlast_lin_sigma')
            return latent_mean, log_latent_sigmas
        else:
            return ops.linear(opts, layer_x, opts['latent_space_dim'], scope='hlast_lin')

    def began_encoder(self, opts, input_, is_training=False, reuse=False, keep_prob=1.):
        num_units = opts['e_num_filters']
        assert num_units == opts['g_num_filters'], 'BEGAN requires same number of filters in encoder and decoder'
        num_layers = opts['e_num_layers']
        layer_x = ops.conv2d(opts, input_, num_units, scope='h_first_conv')
        for i in xrange(num_layers):
            if i % 3 < 2:
                if i != num_layers - 2:
                    ii = i - (i / 3)
                    scale = (ii + 1 - ii / 2)
                else:
                    ii = i - (i / 3)
                    scale = (ii - (ii - 1) / 2)
                layer_x = ops.conv2d(opts, layer_x, num_units * scale, d_h=1, d_w=1, scope='h%d_conv' % i)
                layer_x = tf.nn.elu(layer_x)
            else:
                if i != num_layers - 1:
                    layer_x = ops.downsample(layer_x, scope='h%d_maxpool' % i, reuse=reuse)
        # Tensor should be [N, 8, 8, filters] right now

        if opts['e_is_random']:
            latent_mean = ops.linear(
                opts, layer_x, opts['latent_space_dim'], scope='hlast_lin')
            log_latent_sigmas = ops.linear(
                opts, layer_x, opts['latent_space_dim'], scope='hlast_lin_sigma')
            return latent_mean, log_latent_sigmas
        else:
            return ops.linear(opts, layer_x, opts['latent_space_dim'], scope='hlast_lin')

    def _data_augmentation(self, opts, real_points, is_training):
        if not opts['data_augm']:
            return real_points

        height = int(real_points.get_shape()[1])
        width = int(real_points.get_shape()[2])
        depth = int(real_points.get_shape()[3])
        # logging.error("real_points shape", real_points.get_shape())
        def _distort_func(image):
            # tf.image.per_image_standardization(image), should we?
            # Pad with zeros.
            image = tf.image.resize_image_with_crop_or_pad(
                image, height+4, width+4)
            image = tf.random_crop(image, [height, width, depth])
            image = tf.image.random_flip_left_right(image)
            image = tf.image.random_brightness(image, max_delta=0.1)
            image = tf.minimum(tf.maximum(image, 0.0), 1.0)
            image = tf.image.random_contrast(image, lower=0.8, upper=1.3)
            image = tf.minimum(tf.maximum(image, 0.0), 1.0)
            image = tf.image.random_hue(image, 0.08)
            image = tf.minimum(tf.maximum(image, 0.0), 1.0)
            image = tf.image.random_saturation(image, lower=0.8, upper=1.3)
            image = tf.minimum(tf.maximum(image, 0.0), 1.0)
            return image

        def _regular_func(image):
            # tf.image.per_image_standardization(image)?
            return image

        distorted_images = tf.cond(
            is_training,
            lambda: tf.map_fn(_distort_func, real_points,
                              parallel_iterations=100),
            lambda: tf.map_fn(_regular_func, real_points,
                              parallel_iterations=100))

        return distorted_images

    def _recon_loss_using_disc_encoder(
            self, opts, reconstructed_training, encoded_training,
            real_points, is_training_ph, keep_prob_ph):
        """Build an additional loss using the encoder as discriminator."""
        reconstructed_reencoded_sg = self.encoder(
            opts, tf.stop_gradient(reconstructed_training),
            is_training=is_training_ph, keep_prob=keep_prob_ph, reuse=True)
        if opts['e_is_random']:
            reconstructed_reencoded_sg = reconstructed_reencoded_sg[0]
        reconstructed_reencoded = self.encoder(
            opts, reconstructed_training, is_training=is_training_ph,
            keep_prob=keep_prob_ph, reuse=True)
        if opts['e_is_random']:
            reconstructed_reencoded = reconstructed_reencoded[0]
        # Below line enforces the forward to be reconstructed_reencoded and backwards to NOT change the encoder....
        crazy_hack = reconstructed_reencoded - reconstructed_reencoded_sg +\
            tf.stop_gradient(reconstructed_reencoded_sg)
        encoded_training_sg = self.encoder(
            opts, tf.stop_gradient(real_points),
            is_training=is_training_ph, keep_prob=keep_prob_ph, reuse=True)
        if opts['e_is_random']:
            encoded_training_sg = encoded_training_sg[0]

        adv_fake_layer = ops.linear(opts, reconstructed_reencoded_sg, 1, scope='adv_layer')
        adv_true_layer = ops.linear(opts, encoded_training_sg, 1, scope='adv_layer', reuse=True)
        adv_fake = tf.nn.sigmoid_cross_entropy_with_logits(
                    logits=adv_fake_layer, labels=tf.zeros_like(adv_fake_layer))
        adv_true = tf.nn.sigmoid_cross_entropy_with_logits(
                    logits=adv_true_layer, labels=tf.ones_like(adv_true_layer))
        adv_fake = tf.reduce_mean(adv_fake)
        adv_true = tf.reduce_mean(adv_true)
        adv_c_loss = adv_fake + adv_true
        emb_c = tf.reduce_sum(tf.square(crazy_hack - tf.stop_gradient(encoded_training)), 1)
        emb_c_loss = tf.reduce_mean(tf.sqrt(emb_c + 1e-5))
        # Normalize the loss, so that it does not depend on how good the
        # discriminator is.
        emb_c_loss = emb_c_loss / tf.stop_gradient(emb_c_loss)
        return adv_c_loss, emb_c_loss

    def _recon_loss_using_disc_conv(self, opts, reconstructed_training, real_points, is_training, keep_prob):
        """Build an additional loss using a discriminator in X space."""
        def _conv_flatten(x, kernel_size):
            height = int(x.get_shape()[1])
            width = int(x.get_shape()[2])
            channels = int(x.get_shape()[3])
            w_sum = tf.eye(num_rows=channels, num_columns=channels, batch_shape=[kernel_size * kernel_size])
            w_sum = tf.reshape(w_sum, [kernel_size, kernel_size, channels, channels])
            w_sum = w_sum / (kernel_size * kernel_size)
            sum_ = tf.nn.conv2d(x, w_sum, strides=[1, 1, 1, 1], padding='SAME')
            size = prod_dim(sum_)
            assert size == height * width * channels, size
            return tf.reshape(sum_, [-1, size])

        def _gram_scores(tensor, kernel_size):
            assert len(tensor.get_shape()) == 4, tensor
            ttensor = tf.transpose(tensor, [3, 1, 2, 0])
            rand_indices = tf.random_shuffle(tf.range(ttensor.get_shape()[0]))
            shuffled = tf.gather(ttensor, rand_indices)

            shuffled = tf.transpose(shuffled, [3, 1, 2, 0])
            cross_p = _conv_flatten(tensor * shuffled, kernel_size)  # shape [batch_size, height * width * channels]
            diag_p = _conv_flatten(tf.square(tensor), kernel_size)  # shape [batch_size, height * width * channels]
            return cross_p, diag_p

        def _architecture(inputs, reuse=None):
            with tf.variable_scope('DISC_X_LOSS', reuse=reuse):
                num_units = opts['adv_c_num_units']
                num_layers = 1
                filter_sizes = opts['adv_c_patches_size']
                if isinstance(filter_sizes, int):
                    filter_sizes = [filter_sizes]
                else:
                    filter_sizes = [int(n) for n in filter_sizes.split(',')]
                embedded_outputs = []
                linear_outputs = []
                for filter_size in filter_sizes:
                    layer_x = inputs
                    for i in xrange(num_layers):
    #                     scale = 2**(num_layers-i-1)
                        layer_x = ops.conv2d(opts, layer_x, num_units, d_h=1, d_w=1, scope='h%d_conv%d' % (i, filter_size),
                                             conv_filters_dim=filter_size, padding='SAME')
    #                     if opts['batch_norm']:
    #                         layer_x = ops.batch_norm(opts, layer_x, is_training, reuse, scope='bn%d_%d' % (i, filter_size))
                        layer_x = ops.lrelu(layer_x, 0.1)
                    last = ops.conv2d(
                        opts, layer_x, 1, d_h=1, d_w=1, scope="last_lin%d" % filter_size, conv_filters_dim=1, l2_norm=True)
                    if opts['cross_p_w'] > 0.0 or opts['diag_p_w'] > 0.0:
                        cross_p, diag_p = _gram_scores(layer_x, filter_size)
                        embedded_outputs.append(cross_p * opts['cross_p_w'])
                        embedded_outputs.append(diag_p * opts['diag_p_w'])
                    fl = flatten(layer_x)
#                     fl = tf.Print(fl, [fl], "fl")
                    embedded_outputs.append(fl)
                    size = int(last.get_shape()[1])
                    linear_outputs.append(tf.reshape(last, [-1, size * size]))
                if len(embedded_outputs) > 1:
                    embedded_outputs = tf.concat(embedded_outputs, 1)
                else:
                    embedded_outputs = embedded_outputs[0]
                if len(linear_outputs) > 1:
                    linear_outputs = tf.concat(linear_outputs, 1)
                else:
                    linear_outputs = linear_outputs[0]

                return embedded_outputs, linear_outputs


        reconstructed_embed_sg, adv_fake_layer = _architecture(tf.stop_gradient(reconstructed_training), reuse=None)
        reconstructed_embed, _ = _architecture(reconstructed_training, reuse=True)
        # Below line enforces the forward to be reconstructed_embed and backwards to NOT change the discriminator....
        crazy_hack = reconstructed_embed-reconstructed_embed_sg+tf.stop_gradient(reconstructed_embed_sg)
        real_p_embed_sg, adv_true_layer = _architecture(tf.stop_gradient(real_points), reuse=True)
        real_p_embed, _ = _architecture(real_points, reuse=True)

        adv_fake = tf.nn.sigmoid_cross_entropy_with_logits(
                    logits=adv_fake_layer, labels=tf.zeros_like(adv_fake_layer))
        adv_true = tf.nn.sigmoid_cross_entropy_with_logits(
                    logits=adv_true_layer, labels=tf.ones_like(adv_true_layer))
        adv_fake = tf.reduce_mean(adv_fake)
        adv_true = tf.reduce_mean(adv_true)

        adv_c_loss = adv_fake + adv_true
        emb_c = tf.reduce_mean(tf.square(crazy_hack - tf.stop_gradient(real_p_embed)), 1)

        real_points_shuffle = tf.stop_gradient(tf.random_shuffle(real_p_embed))
        emb_c_shuffle = tf.reduce_mean(tf.square(real_points_shuffle - tf.stop_gradient(reconstructed_embed)), 1)

        raw_emb_c_loss = tf.reduce_mean(emb_c)
        shuffled_emb_c_loss = tf.reduce_mean(emb_c_shuffle)
        emb_c_loss = raw_emb_c_loss / shuffled_emb_c_loss
        emb_c_loss = emb_c_loss * 40

        return adv_c_loss, emb_c_loss

    def _recon_loss_using_disc_conv_eb(self, opts, reconstructed_training, real_points, is_training, keep_prob):
        """Build an additional loss using a discriminator in X space, using Energy Based approach."""
        def copy3D(height, width, channels):
            m = np.zeros([height, width, channels, height, width, channels])
            for i in xrange(height):
                for j in xrange(width):
                    for c in xrange(channels):
                        m[i, j, c, i, j, c] = 1.0
            return tf.constant(np.reshape(m, [height, width, channels, -1]), dtype=tf.float32)

        def _architecture(inputs, reuse=None):
            dim = opts['adv_c_patches_size']
            height = int(inputs.get_shape()[1])
            width = int(inputs.get_shape()[2])
            channels = int(inputs.get_shape()[3])
            with tf.variable_scope('DISC_X_LOSS', reuse=reuse):
                num_units = opts['adv_c_num_units']
                num_layers = 1
                layer_x = inputs
                for i in xrange(num_layers):
#                     scale = 2**(num_layers-i-1)
                    layer_x = ops.conv2d(opts, layer_x, num_units, d_h=1, d_w=1, scope='h%d_conv' % i,
                                         conv_filters_dim=dim, padding='SAME')
#                     if opts['batch_norm']:
#                         layer_x = ops.batch_norm(opts, layer_x, is_training, reuse, scope='bn%d' % i)
                    layer_x = ops.lrelu(layer_x, 0.1)  #tf.nn.relu(layer_x)

                copy_w = copy3D(dim, dim, channels)
                duplicated = tf.nn.conv2d(inputs, copy_w, strides=[1, 1, 1, 1], padding='SAME')
                decoded = ops.conv2d(
                    opts, layer_x, channels * dim * dim, d_h=1, d_w=1, scope="decoder",
                    conv_filters_dim=1, padding='SAME')
            reconstruction = tf.reduce_mean(tf.square(tf.stop_gradient(duplicated) - decoded), [1, 2, 3])
            assert len(reconstruction.get_shape()) == 1
            return flatten(layer_x), reconstruction


        reconstructed_embed_sg, adv_fake_layer = _architecture(tf.stop_gradient(reconstructed_training), reuse=None)
        reconstructed_embed, _ = _architecture(reconstructed_training, reuse=True)
        # Below line enforces the forward to be reconstructed_embed and backwards to NOT change the discriminator....
        crazy_hack = reconstructed_embed-reconstructed_embed_sg+tf.stop_gradient(reconstructed_embed_sg)
        real_p_embed_sg, adv_true_layer = _architecture(tf.stop_gradient(real_points), reuse=True)
        real_p_embed, _ = _architecture(real_points, reuse=True)

        adv_fake = tf.reduce_mean(adv_fake_layer)
        adv_true = tf.reduce_mean(adv_true_layer)

        adv_c_loss = tf.log(adv_true) - tf.log(adv_fake)
        emb_c = tf.reduce_sum(tf.square(crazy_hack - tf.stop_gradient(real_p_embed)), 1)
        emb_c_loss = tf.reduce_mean(emb_c)

        return adv_c_loss, emb_c_loss

    def _recon_loss_using_vgg(self, opts, reconstructed_training, real_points, is_training, keep_prob):
        """Build an additional loss using a pretrained VGG in X space."""

        def _architecture(_inputs, reuse=None):
            _, end_points = vgg_16(_inputs, is_training=is_training, dropout_keep_prob=keep_prob, reuse=reuse)
            layer_name = opts['vgg_layer']
            if layer_name == 'concat':
                outputs = []
                for ln in ['pool1', 'pool2', 'pool3']:
                    output = end_points[ln]
                    output = flatten(output)
                    outputs.append(output)
                output = tf.concat(outputs, 1)
            elif layer_name.startswith('concat_w'):
                weights = layer_name.split(',')[1:]
                assert len(weights) == 5
                outputs = []
                for lnum in range(5):
                    num = lnum + 1
                    ln = 'pool%d' % num
                    output = end_points[ln]
                    output = flatten(output)
                    # We sqrt the weight here because we use L2 after.
                    outputs.append(np.sqrt(float(weights[lnum])) * output)
                output = tf.concat(outputs, 1)
            else:
                output = end_points[layer_name]
                output = flatten(output)
            if reuse is None:
                variables_to_restore = slim.get_variables_to_restore(include=['vgg_16'])
                path = os.path.join(opts['data_dir'], 'vgg_16.ckpt')
#                 '/tmpp/models/vgg_16.ckpt'
                init_assign_op, init_feed_dict = slim.assign_from_checkpoint(path, variables_to_restore)
                self._additional_init_ops += [init_assign_op]
                self._init_feed_dict.update(init_feed_dict)
            return output


        reconstructed_embed_sg = _architecture(tf.stop_gradient(reconstructed_training), reuse=None)
        reconstructed_embed = _architecture(reconstructed_training, reuse=True)
        # Below line enforces the forward to be reconstructed_embed and backwards to NOT change the discriminator....
        crazy_hack = reconstructed_embed-reconstructed_embed_sg+tf.stop_gradient(reconstructed_embed_sg)
        real_p_embed = _architecture(real_points, reuse=True)

        emb_c = tf.reduce_mean(tf.square(crazy_hack - tf.stop_gradient(real_p_embed)), 1)
        emb_c_loss = tf.reduce_mean(tf.sqrt(emb_c + 1e-5))
#         emb_c_loss = tf.Print(emb_c_loss, [emb_c_loss], "emb_c_loss")
#         # Normalize the loss, so that it does not depend on how good the
#         # discriminator is.
#         emb_c_loss = emb_c_loss / tf.stop_gradient(emb_c_loss)
        return emb_c_loss

    def _recon_loss_using_moments(self, opts, reconstructed_training, real_points, is_training, keep_prob):
        """Build an additional loss using moments."""

        def _architecture(_inputs):
            return compute_moments(_inputs, moments=[2])  # TODO

        reconstructed_embed = _architecture(reconstructed_training)
        real_p_embed = _architecture(real_points)

        emb_c = tf.reduce_mean(tf.square(reconstructed_embed - tf.stop_gradient(real_p_embed)), 1)
#         emb_c = tf.Print(emb_c, [emb_c], "emb_c")
        emb_c_loss = tf.reduce_mean(emb_c)
        return emb_c_loss * 100.0 * 100.0 # TODO: constant.

    def _recon_loss_using_vgg_moments(self, opts, reconstructed_training, real_points, is_training, keep_prob):
        """Build an additional loss using a pretrained VGG in X space."""

        def _architecture(_inputs, reuse=None):
            _, end_points = vgg_16(_inputs, is_training=is_training, dropout_keep_prob=keep_prob, reuse=reuse)
            layer_name = opts['vgg_layer']
            output = end_points[layer_name]
#             output = flatten(output)
            output /= 255.0  # the vgg_16 method scales everything by 255.0, so we divide back here.
            variances = compute_moments(output, moments=[2])

            if reuse is None:
                variables_to_restore = slim.get_variables_to_restore(include=['vgg_16'])
                path = os.path.join(opts['data_dir'], 'vgg_16.ckpt')
#                 '/tmpp/models/vgg_16.ckpt'
                init_assign_op, init_feed_dict = slim.assign_from_checkpoint(path, variables_to_restore)
                self._additional_init_ops += [init_assign_op]
                self._init_feed_dict.update(init_feed_dict)
            return variances

        reconstructed_embed_sg = _architecture(tf.stop_gradient(reconstructed_training), reuse=None)
        reconstructed_embed = _architecture(reconstructed_training, reuse=True)
        # Below line enforces the forward to be reconstructed_embed and backwards to NOT change the discriminator....
        crazy_hack = reconstructed_embed-reconstructed_embed_sg+tf.stop_gradient(reconstructed_embed_sg)
        real_p_embed = _architecture(real_points, reuse=True)

        emb_c = tf.reduce_mean(tf.square(crazy_hack - tf.stop_gradient(real_p_embed)), 1)
        emb_c_loss = tf.reduce_mean(emb_c)
#         emb_c_loss = tf.Print(emb_c_loss, [emb_c_loss], "emb_c_loss")
#         # Normalize the loss, so that it does not depend on how good the
#         # discriminator is.
#         emb_c_loss = emb_c_loss / tf.stop_gradient(emb_c_loss)
        return emb_c_loss   # TODO: constant.

    def add_least_gaussian2d_ops(self, opts):
        """ Add ops searching for the 2d plane in z_dim hidden space
            corresponding to the 'least Gaussian' look of the sample
        """

        with tf.variable_scope('leastGaussian2d'):
            # Projection matrix which we are going to tune
            sample_ph = tf.placeholder(
                tf.float32, [None, opts['latent_space_dim']],
                name='sample_ph')
            v = tf.get_variable(
                "proj_v", [opts['latent_space_dim'], 1],
                tf.float32, tf.random_normal_initializer(stddev=1.))
            u = tf.get_variable(
                "proj_u", [opts['latent_space_dim'], 1],
                tf.float32, tf.random_normal_initializer(stddev=1.))
        npoints = tf.cast(tf.shape(sample_ph)[0], tf.int32)
        # First we need to make sure projection matrix is orthogonal
        v_norm = tf.nn.l2_normalize(v, 0)
        dotprod = tf.reduce_sum(tf.multiply(u, v_norm))
        u_ort = u - dotprod * v_norm
        u_norm = tf.nn.l2_normalize(u_ort, 0)
        Mproj = tf.concat([v_norm, u_norm], 1)
        sample_proj = tf.matmul(sample_ph, Mproj)
        a = tf.eye(npoints) - tf.ones([npoints, npoints]) / tf.cast(npoints, tf.float32)
        b = tf.matmul(sample_proj, tf.matmul(a, a), transpose_a=True)
        b = tf.matmul(b, sample_proj)
        # Sample covariance matrix
        covhat = b / (tf.cast(npoints, tf.float32) - 1)
        # covhat = tf.Print(covhat, [covhat], 'Cov:')
        with tf.variable_scope('leastGaussian2d'):
            gcov = opts['pot_pz_std'] * opts['pot_pz_std'] * tf.eye(2)
            # l2 distance between sample cov and the Gaussian cov
            projloss =  tf.reduce_sum(tf.square(covhat - gcov))
            # Also account for the first moment, i.e. expected value
            projloss += tf.reduce_sum(tf.square(tf.reduce_mean(sample_proj, 0)))
            # We are maximizing
            projloss = -projloss
            optim = tf.train.AdamOptimizer(0.001, 0.9)
            optim = optim.minimize(projloss, var_list=[v, u])

        self._proj_u = u_norm
        self._proj_v = v_norm
        self._proj_sample_ph = sample_ph
        self._proj_covhat = covhat
        self._proj_loss = projloss
        self._proj_optim = optim

    def least_gaussian_2d(self, opts, X):
        """
        Given a sample X of shape (n_points, n_z) find 2d plain
        such that projection looks least gaussian.
        """
        with self._session.as_default(), self._session.graph.as_default():
            sample_ph = self._proj_sample_ph
            optim = self._proj_optim
            loss = self._proj_loss
            u = self._proj_u
            v = self._proj_v
            covhat = self._proj_covhat
            proj_mat = tf.concat([v, u], 1).eval()
            dot_prod = -1
            best_of_runs = 10e5 # Any positive value would do
            updated = False
            for _start in xrange(3):
                # We will run 3 times from random inits
                loss_prev = 10e5 # Any positive value would do
                proj_vars = tf.get_collection(
                    tf.GraphKeys.GLOBAL_VARIABLES, scope="leastGaussian2d")
                self._session.run(tf.variables_initializer(proj_vars))
                step = 0
                for _ in xrange(5000):
                    self._session.run(optim, feed_dict={sample_ph:X})
                    step += 1
                    if step % 10 == 0:
                        loss_cur = loss.eval(feed_dict={sample_ph: X})
                        rel_imp = abs(loss_cur - loss_prev) / abs(loss_prev)
                        if rel_imp < 1e-2:
                            break
                        loss_prev = loss_cur
                loss_final = loss.eval(feed_dict={sample_ph: X})
                if loss_final < best_of_runs:
                    updated = True
                    best_of_runs = loss_final
                    proj_mat = tf.concat([v, u], 1).eval()
                    dot_prod = tf.reduce_sum(tf.multiply(u, v)).eval()
        if not updated:
            logging.error('WARNING: possible bug in the worst 2d projection')
        return proj_mat, dot_prod

    def _build_model_internal(self, opts):
        """Build the Graph corresponding to POT implementation.

        """
        data_shape = self._data.data_shape
        additional_losses = collections.OrderedDict()

        # Placeholders
        real_points_ph = tf.placeholder(
            tf.float32, [None] + list(data_shape), name='real_points_ph')
        noise_ph = tf.placeholder(
            tf.float32, [None] + [opts['latent_space_dim']], name='noise_ph')
        enc_noise_ph = tf.placeholder(
            tf.float32, [None] + [opts['latent_space_dim']], name='enc_noise_ph')
        lr_decay_ph = tf.placeholder(tf.float32)
        is_training_ph = tf.placeholder(tf.bool, name='is_training_ph')
        keep_prob_ph = tf.placeholder(tf.float32, name='keep_prob_ph')

        # Operations
        if opts['pz_transform']:
            assert opts['z_test'] == 'gan', 'Pz transforms are currently allowed only for POT+GAN'
            noise = self.pz_sampler(opts, noise_ph)
        else:
            noise = noise_ph

        real_points = self._data_augmentation(
            opts, real_points_ph, is_training_ph)

        if opts['e_is_random']:
            # If encoder is random we map the training points
            # to the expectation of Q(Z|X) and then add the scaled
            # Gaussian noise corresponding to the learned sigmas
            enc_train_mean, enc_log_sigmas = self.encoder(
                opts, real_points,
                is_training=is_training_ph, keep_prob=keep_prob_ph)
            # enc_log_sigmas = tf.Print(enc_log_sigmas, [tf.reduce_max(enc_log_sigmas),
            #                                            tf.reduce_min(enc_log_sigmas),
            #                                            tf.reduce_mean(enc_log_sigmas)], 'Log sigmas:')
            # enc_log_sigmas = tf.Print(enc_log_sigmas, [tf.slice(enc_log_sigmas, [0,0], [1,-1])], 'Log sigmas:')
            # stds = tf.sqrt(tf.exp(enc_log_sigmas) + 1e-05)
            stds = tf.sqrt(tf.nn.relu(enc_log_sigmas) + 1e-05)
            # stds = tf.Print(stds, [stds[0], stds[1], stds[2], stds[3]], 'Stds: ')
            # stds = tf.Print(stds, [enc_train_mean[0], enc_train_mean[1], enc_train_mean[2]], 'Means: ')
            scaled_noise = tf.multiply(stds, enc_noise_ph)
            encoded_training = enc_train_mean + scaled_noise
        else:
            encoded_training = self.encoder(
                opts, real_points,
                is_training=is_training_ph, keep_prob=keep_prob_ph)
        reconstructed_training = self.generator(
            opts, encoded_training,
            is_training=is_training_ph, keep_prob=keep_prob_ph)
        reconstructed_training.set_shape(real_points.get_shape())

        if opts['recon_loss'] == 'l2':
            # c(x,y) = ||x - y||_2
            loss_reconstr = tf.reduce_sum(
                tf.square(real_points - reconstructed_training), axis=1)
            # sqrt(x + delta) guarantees the direvative 1/(x + delta) is finite
            loss_reconstr = tf.reduce_mean(tf.sqrt(loss_reconstr + 1e-08))
        elif opts['recon_loss'] == 'l2f':
            # c(x,y) = ||x - y||_2
            loss_reconstr = tf.reduce_sum(
                tf.square(real_points - reconstructed_training), axis=[1, 2, 3])
            loss_reconstr = tf.reduce_mean(tf.sqrt(1e-08 + loss_reconstr)) * 0.2
        elif opts['recon_loss'] == 'l2sq':
            # c(x,y) = ||x - y||_2^2
            loss_reconstr = tf.reduce_sum(
                tf.square(real_points - reconstructed_training), axis=[1, 2, 3])
            loss_reconstr = tf.reduce_mean(loss_reconstr) * 0.05
        elif opts['recon_loss'] == 'l1':
            # c(x,y) = ||x - y||_1
            loss_reconstr = tf.reduce_mean(tf.reduce_sum(
                tf.abs(real_points - reconstructed_training), axis=[1, 2, 3])) * 0.02
        else:
            assert False

        # Pearson independence test of coordinates in Z space
        loss_z_corr = self.correlation_loss(opts, encoded_training)
        # Perform a Qz = Pz goodness of fit test based on Stein Discrepancy
        if opts['z_test'] == 'gan':
            # Pz = Qz test based on GAN in the Z space
            d_logits_Pz = self.discriminator(opts, noise)
            d_logits_Qz = self.discriminator(opts, encoded_training, reuse=True)
            d_loss_Pz = tf.reduce_mean(
                tf.nn.sigmoid_cross_entropy_with_logits(
                    logits=d_logits_Pz, labels=tf.ones_like(d_logits_Pz)))
            d_loss_Qz = tf.reduce_mean(
                tf.nn.sigmoid_cross_entropy_with_logits(
                    logits=d_logits_Qz, labels=tf.zeros_like(d_logits_Qz)))
            d_loss_Qz_trick = tf.reduce_mean(
                tf.nn.sigmoid_cross_entropy_with_logits(
                    logits=d_logits_Qz, labels=tf.ones_like(d_logits_Qz)))
            d_loss = opts['pot_lambda'] * (d_loss_Pz + d_loss_Qz)
            if opts['pz_transform']:
                loss_match = d_loss_Qz_trick - d_loss_Pz
            else:
                loss_match = d_loss_Qz_trick
        elif opts['z_test'] == 'mmd':
            # Pz = Qz test based on MMD(Pz, Qz)
            loss_match = self.discriminator_mmd_test(opts, encoded_training, noise)
            d_loss = None
            d_logits_Pz = None
            d_logits_Qz = None
        elif opts['z_test'] == 'lks':
            # Pz = Qz test without adversarial training
            # based on Kernel Stein Discrepancy
            # Uncomment next line to check for the real Pz
            # loss_match = self.discriminator_test(opts, noise_ph)
            loss_match = self.discriminator_test(opts, encoded_training)
            d_loss = None
            d_logits_Pz = None
            d_logits_Qz = None
        else:
            # Pz = Qz test without adversarial training
            # (a) Check for multivariate Gaussianity
            #     by checking Gaussianity of all the 1d projections
            # (b) Run Pearson's test of coordinate independance
            loss_match = self.discriminator_test(opts, encoded_training)
            loss_match = loss_match + opts['z_test_corr_w'] * loss_z_corr
            d_loss = None
            d_logits_Pz = None
            d_logits_Qz = None
        g_mom_stats = self.moments_stats(opts, encoded_training)
        loss = opts['reconstr_w'] * loss_reconstr + opts['pot_lambda'] * loss_match

        # Optionally, add one more cost function based on the embeddings
        # add a discriminator in the X space, reusing the encoder or a new model.
        if opts['adv_c_loss'] == 'encoder':
            adv_c_loss, emb_c_loss = self._recon_loss_using_disc_encoder(
                opts, reconstructed_training, encoded_training, real_points, is_training_ph, keep_prob_ph)
            loss += opts['adv_c_loss_w'] * adv_c_loss + opts['emb_c_loss_w'] * emb_c_loss
            additional_losses['adv_c'], additional_losses['emb_c'] = adv_c_loss, emb_c_loss
        elif opts['adv_c_loss'] == 'conv':
            adv_c_loss, emb_c_loss = self._recon_loss_using_disc_conv(
                opts, reconstructed_training, real_points, is_training_ph, keep_prob_ph)
            additional_losses['adv_c'], additional_losses['emb_c'] = adv_c_loss, emb_c_loss
            loss += opts['adv_c_loss_w'] * adv_c_loss + opts['emb_c_loss_w'] * emb_c_loss
        elif opts['adv_c_loss'] == 'conv_eb':
            adv_c_loss, emb_c_loss = self._recon_loss_using_disc_conv_eb(
                opts, reconstructed_training, real_points, is_training_ph, keep_prob_ph)
            additional_losses['adv_c'], additional_losses['emb_c'] = adv_c_loss, emb_c_loss
            loss += opts['adv_c_loss_w'] * adv_c_loss + opts['emb_c_loss_w'] * emb_c_loss
        elif opts['adv_c_loss'] == 'vgg':
            emb_c_loss = self._recon_loss_using_vgg(
                opts, reconstructed_training, real_points, is_training_ph, keep_prob_ph)
            loss += opts['emb_c_loss_w'] * emb_c_loss
            additional_losses['emb_c'] = emb_c_loss
        elif opts['adv_c_loss'] == 'moments':
            emb_c_loss = self._recon_loss_using_moments(
                opts, reconstructed_training, real_points, is_training_ph, keep_prob_ph)
            loss += opts['emb_c_loss_w'] * emb_c_loss
            additional_losses['emb_c'] = emb_c_loss
        elif opts['adv_c_loss'] == 'vgg_moments':
            emb_c_loss = self._recon_loss_using_vgg_moments(
                opts, reconstructed_training, real_points, is_training_ph, keep_prob_ph)
            loss += opts['emb_c_loss_w'] * emb_c_loss
            additional_losses['emb_c'] = emb_c_loss
        else:
            assert opts['adv_c_loss'] == 'none'

        # Add ops to pretrain the Qz match mean and covariance of Pz
        loss_pretrain = None
        if opts['e_pretrain']:
            # Next two vectors are zdim-dimensional
            mean_pz = tf.reduce_mean(noise, axis=0, keep_dims=True)
            mean_qz = tf.reduce_mean(encoded_training, axis=0, keep_dims=True)
            mean_loss = tf.reduce_mean(tf.square(mean_pz - mean_qz))
            cov_pz = tf.matmul(noise - mean_pz,
                               noise - mean_pz, transpose_a=True)
            cov_pz /= opts['e_pretrain_bsize'] - 1.
            cov_qz = tf.matmul(encoded_training - mean_qz,
                               encoded_training - mean_qz, transpose_a=True)
            cov_qz /= opts['e_pretrain_bsize'] - 1.
            cov_loss = tf.reduce_mean(tf.square(cov_pz - cov_qz))
            loss_pretrain = mean_loss + cov_loss

        # Also add ops to find the least Gaussian 2d projection 
        # this is handy when visually inspection Qz = Pz
        self.add_least_gaussian2d_ops(opts)

        # Optimizer ops
        t_vars = tf.trainable_variables()
        # Updates for discriminator
        d_vars = [var for var in t_vars if 'DISCRIMINATOR/' in var.name]
        # Updates for everything but adversary (encoder, decoder and possibly pz-transform)
        all_vars = [var for var in t_vars if 'DISCRIMINATOR/' not in var.name]
        # Updates for everything but adversary (encoder, decoder and possibly pz-transform)
        eg_vars = [var for var in t_vars if 'GENERATOR/' in var.name or 'ENCODER/' in var.name]
        # Encoder variables separately if we want to pretrain
        e_vars = [var for var in t_vars if 'ENCODER/' in var.name]

        logging.error('Param num in G and E: %d' % \
                np.sum([np.prod([int(d) for d in v.get_shape()]) for v in eg_vars]))
        for v in eg_vars:
            print v.name, [int(d) for d in v.get_shape()]

        if len(d_vars) > 0:
            d_optim = ops.optimizer(opts, net='d', decay=lr_decay_ph).minimize(loss=d_loss, var_list=d_vars)
        else:
            d_optim = None
        optim = ops.optimizer(opts, net='g', decay=lr_decay_ph).minimize(loss=loss, var_list=all_vars)
        pretrain_optim = None
        if opts['e_pretrain']:
            pretrain_optim = ops.optimizer(opts, net='g').minimize(loss=loss_pretrain, var_list=e_vars)


        generated_images = self.generator(
            opts, noise, is_training=is_training_ph,
            reuse=True, keep_prob=keep_prob_ph)

        self._real_points_ph = real_points_ph
        self._real_points = real_points
        self._noise_ph = noise_ph
        self._noise = noise
        self._enc_noise_ph = enc_noise_ph
        self._lr_decay_ph = lr_decay_ph
        self._is_training_ph = is_training_ph
        self._keep_prob_ph = keep_prob_ph
        self._optim = optim
        self._d_optim = d_optim
        self._pretrain_optim = pretrain_optim
        self._loss = loss
        self._loss_reconstruct = loss_reconstr
        self._loss_match = loss_match
        self._loss_z_corr = loss_z_corr
        self._loss_pretrain = loss_pretrain
        self._additional_losses = additional_losses
        self._g_mom_stats = g_mom_stats
        self._d_loss = d_loss
        self._generated = generated_images
        self._Qz = encoded_training
        self._reconstruct_x = reconstructed_training

        saver = tf.train.Saver(max_to_keep=10)
        tf.add_to_collection('real_points_ph', self._real_points_ph)
        tf.add_to_collection('noise_ph', self._noise_ph)
        tf.add_to_collection('enc_noise_ph', self._enc_noise_ph)
        if opts['pz_transform']:
            tf.add_to_collection('noise', self._noise)
        tf.add_to_collection('is_training_ph', self._is_training_ph)
        tf.add_to_collection('keep_prob_ph', self._keep_prob_ph)
        tf.add_to_collection('encoder', self._Qz)
        tf.add_to_collection('decoder', self._generated)
        if d_logits_Pz is not None:
            tf.add_to_collection('disc_logits_Pz', d_logits_Pz)
        if d_logits_Qz is not None:
            tf.add_to_collection('disc_logits_Qz', d_logits_Qz)

        self._saver = saver

        logging.error("Building Graph Done.")

    def pretrain(self, opts):
        steps_max = 200
        batch_size = opts['e_pretrain_bsize']
        for step in xrange(steps_max):
            train_size = self._data.num_points
            data_ids = np.random.choice(train_size, min(train_size, batch_size),
                                        replace=False)
            batch_images = self._data.data[data_ids].astype(np.float)
            batch_noise = opts['pot_pz_std'] *\
                utils.generate_noise(opts, batch_size)
            # Noise for the random encoder (if present)
            batch_enc_noise = utils.generate_noise(opts, batch_size)

            # Update encoder
            [_, loss_pretrain] = self._session.run(
                [self._pretrain_optim,
                 self._loss_pretrain],
                feed_dict={self._real_points_ph: batch_images,
                           self._noise_ph: batch_noise,
                           self._enc_noise_ph: batch_enc_noise,
                           self._is_training_ph: True,
                           self._keep_prob_ph: opts['dropout_keep_prob']})

            if opts['verbose'] == 2:
                logging.error('Step %d/%d, loss=%f' % (step, steps_max, loss_pretrain))

            if loss_pretrain < 0.1:
                break

    def _train_internal(self, opts):
        """Train a POT model.

        """
        logging.error(opts)

        batches_num = self._data.num_points / opts['batch_size']
        train_size = self._data.num_points
        num_plot = 320
        sample_prev = np.zeros([num_plot] + list(self._data.data_shape))
        l2s = []
        losses = []
        losses_rec = []
        losses_match = []
        wait = 0

        start_time = time.time()
        counter = 0
        decay = 1.
        logging.error('Training POT')

        # Optionally we first pretrain the Qz to match mean and
        # covariance of Pz
        if opts['e_pretrain']:
            logging.error('Pretraining the encoder')
            self.pretrain(opts)
            logging.error('Pretraining the encoder done')

        for _epoch in xrange(opts["gan_epoch_num"]):

            if opts['decay_schedule'] == "manual":
                if _epoch == 30:
                    decay = decay / 2.
                if _epoch == 50:
                    decay = decay / 5.
                if _epoch == 100:
                    decay = decay / 10.
            elif opts['decay_schedule'] != "plateau":
                assert type(1.0 * opts['decay_schedule']) == float
                decay = 1.0 * 10**(-_epoch / float(opts['decay_schedule']))

            if _epoch > 0 and _epoch % opts['save_every_epoch'] == 0:
                os.path.join(opts['work_dir'], opts['ckpt_dir'])
                self._saver.save(self._session,
                                 os.path.join(opts['work_dir'],
                                              opts['ckpt_dir'],
                                              'trained-pot'),
                                 global_step=counter)

            for _idx in xrange(batches_num):
                data_ids = np.random.choice(train_size, opts['batch_size'],
                                            replace=False, p=self._data_weights)
                batch_images = self._data.data[data_ids].astype(np.float)
                # Noise for the Pz=Qz GAN
                batch_noise = opts['pot_pz_std'] *\
                    utils.generate_noise(opts, opts['batch_size'])
                # Noise for the random encoder (if present)
                batch_enc_noise = utils.generate_noise(opts, opts['batch_size'])

                # Update generator (decoder) and encoder
                [_, loss, loss_rec, loss_match] = self._session.run(
                    [self._optim,
                     self._loss,
                     self._loss_reconstruct,
                     self._loss_match],
                    feed_dict={self._real_points_ph: batch_images,
                               self._noise_ph: batch_noise,
                               self._enc_noise_ph: batch_enc_noise,
                               self._lr_decay_ph: decay,
                               self._is_training_ph: True,
                               self._keep_prob_ph: opts['dropout_keep_prob']})

                if opts['decay_schedule'] == "plateau":
                    # First 30 epochs do nothing
                    if _epoch >= 30:
                        # If no significant progress was made in last 10 epochs
                        # then decrease the learning rate.
                        if loss < min(losses[-20 * batches_num:]):
                            wait = 0
                        else:
                            wait += 1
                        if wait > 10 * batches_num:
                            decay = max(decay  / 1.4, 1e-6)
                            logging.error('Reduction in learning rate: %f' % decay)
                            wait = 0
                losses.append(loss)
                losses_rec.append(loss_rec)
                losses_match.append(loss_match)
                if opts['verbose'] >= 2:
                    # logging.error('loss after %d steps : %f' % (counter, losses[-1]))
                    logging.error('loss match  after %d steps : %f' % (counter, losses_match[-1]))

                # Update discriminator in Z space (if any).
                if self._d_optim is not None:
                    for _st in range(opts['d_steps']):
                        if opts['d_new_minibatch']:
                            d_data_ids = np.random.choice(
                                train_size, opts['batch_size'],
                                replace=False, p=self._data_weights)
                            d_batch_images = self._data.data[data_ids].astype(np.float)
                            d_batch_enc_noise = utils.generate_noise(opts, opts['batch_size'])
                        else:
                            d_batch_images = batch_images
                            d_batch_enc_noise = batch_enc_noise
                        _ = self._session.run(
                            [self._d_optim, self._d_loss],
                            feed_dict={self._real_points_ph: d_batch_images,
                                       self._noise_ph: batch_noise,
                                       self._enc_noise_ph: d_batch_enc_noise,
                                       self._lr_decay_ph: decay,
                                       self._is_training_ph: True,
                                       self._keep_prob_ph: opts['dropout_keep_prob']})
                counter += 1
                now = time.time()

                rec_test = None
                if opts['verbose'] and counter % 500 == 0:
                    # Printing (training and test) loss values
                    test = self._data.test_data[:200]
                    [loss_rec_test, rec_test, g_mom_stats, loss_z_corr, additional_losses] = self._session.run(
                        [self._loss_reconstruct, self._reconstruct_x, self._g_mom_stats, self._loss_z_corr,
                         self._additional_losses],
                        feed_dict={self._real_points_ph: test,
                                   self._enc_noise_ph: utils.generate_noise(opts, len(test)),
                                   self._is_training_ph: False,
                                   self._noise_ph: batch_noise,
                                   self._keep_prob_ph: 1e5})
                    debug_str = 'Epoch: %d/%d, batch:%d/%d, batch/sec:%.2f' % (
                        _epoch+1, opts['gan_epoch_num'], _idx+1,
                        batches_num, float(counter) / (now - start_time))
                    debug_str += '  [L=%.5f, Recon=%.5f, GanL=%.5f, Recon_test=%.5f' % (
                        loss, loss_rec, loss_match, loss_rec_test)
                    debug_str += ',' + ', '.join(
                        ['%s=%.2g' % (k, v) for (k, v) in additional_losses.items()])
                    logging.error(debug_str)
                    if opts['verbose'] >= 2:
                        logging.error(g_mom_stats)
                        logging.error(loss_z_corr)
                    if counter % opts['plot_every'] == 0:
                        # plotting the test images.
                        metrics = Metrics()
                        merged = np.vstack([rec_test[:8 * 10], test[:8 * 10]])
                        r_ptr = 0
                        w_ptr = 0
                        for _ in range(8 * 10):
                            merged[w_ptr] = test[r_ptr]
                            merged[w_ptr + 1] = rec_test[r_ptr]
                            r_ptr += 1
                            w_ptr += 2
                        metrics.make_plots(
                            opts,
                            counter,
                            None,
                            merged,
                            prefix='test_reconstr_e%04d_mb%05d_' % (_epoch, _idx))

                if opts['verbose'] and counter % opts['plot_every'] == 0:
                    # Plotting intermediate results
                    metrics = Metrics()
                    # --Random samples from the model
                    points_to_plot, sample_pz = self._session.run(
                        [self._generated, self._noise],
                        feed_dict={
                            self._noise_ph: self._noise_for_plots[0:num_plot],
                            self._is_training_ph: False,
                            self._keep_prob_ph: 1e5})
                    Qz_num = 320
                    sample_Qz = self._session.run(
                        self._Qz,
                        feed_dict={
                            self._real_points_ph: self._data.data[:Qz_num],
                            self._enc_noise_ph: utils.generate_noise(opts, Qz_num),
                            self._is_training_ph: False,
                            self._keep_prob_ph: 1e5})
                    # Searching least Gaussian 2d projection
                    proj_mat, check = self.least_gaussian_2d(opts, sample_Qz)
                    # Projecting samples from Qz and Pz on this 2d plain
                    metrics.Qz = np.dot(sample_Qz, proj_mat)
                    # metrics.Pz = np.dot(self._noise_for_plots, proj_mat)
                    metrics.Pz = np.dot(sample_pz, proj_mat)
                    if self._data.labels != None:
                        metrics.Qz_labels = self._data.labels[:Qz_num]
                    else:
                        metrics.Qz_labels = None
                    metrics.l2s = losses[:]
                    metrics.losses_match = [opts['pot_lambda'] * el for el in losses_match]
                    metrics.losses_rec = [opts['reconstr_w'] * el for el in losses_rec]
                    to_plot = [points_to_plot, 0 * batch_images[:16], batch_images]
                    if rec_test is not None:
                        to_plot += [0 * batch_images[:16], rec_test[:64]]
                    metrics.make_plots(
                        opts,
                        counter,
                        None,
                        np.vstack(to_plot),
                        prefix='sample_e%04d_mb%05d_' % (_epoch, _idx) if rec_test is None \
                                else 'sample_with_test_e%04d_mb%05d_' % (_epoch, _idx))


                    # --Reconstructions for the train and test points
                    num_real_p = 8 * 10
                    reconstructed, real_p = self._session.run(
                        [self._reconstruct_x, self._real_points],
                        feed_dict={
                            self._real_points_ph: self._data.data[:num_real_p],
                            self._enc_noise_ph: utils.generate_noise(opts, num_real_p),
                            self._is_training_ph: True,
                            self._keep_prob_ph: 1e5})
                    points = real_p
                    merged = np.vstack([reconstructed, points])
                    r_ptr = 0
                    w_ptr = 0
                    for _ in range(8 * 10):
                        merged[w_ptr] = points[r_ptr]
                        merged[w_ptr + 1] = reconstructed[r_ptr]
                        r_ptr += 1
                        w_ptr += 2
                    metrics.make_plots(
                        opts,
                        counter,
                        None,
                        merged,
                        prefix='reconstr_e%04d_mb%05d_' % (_epoch, _idx))
                    sample_prev = points_to_plot[:]
        if _epoch > 0:
            os.path.join(opts['work_dir'], opts['ckpt_dir'])
            self._saver.save(self._session,
                             os.path.join(opts['work_dir'],
                                          opts['ckpt_dir'],
                                          'trained-pot-final'),
                             global_step=counter)

    def _sample_internal(self, opts, num):
        """Sample from the trained GAN model.

        """
        # noise = opts['pot_pz_std'] * utils.generate_noise(opts, num)
        # sample = self._run_batch(
        #     opts, self._generated, self._noise_ph, noise, self._is_training_ph, False)
        sample = None
        return sample