Skip to content
Permalink
Browse files

Add files.

  • Loading branch information...
thjashin committed Jun 20, 2018
1 parent 5f19e67 commit 10db359bb95b7b6f9de2e1b0ef2c93d7bd863bde
No changes.
@@ -0,0 +1,7 @@
#!/usr/bin/env python
# -*- coding: utf-8 -*-

from .base import *
from .stein import *
from .spectral import *
from .entropy import *
@@ -0,0 +1,79 @@
#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf


class ScoreEstimator(object):
def __init__(self):
pass

def rbf_kernel(self, x1, x2, kernel_width):
# square1 = tf.expand_dims(tf.reduce_sum(x1**2, -1), -1)
# square2 = tf.expand_dims(tf.reduce_sum(x2**2, -1), -2)
# len_x2 = len(x2.get_shape())
# cross_term = tf.matmul(
# x1, tf.transpose(x2, range(len_x2-2)+[len_x2-1, len_x2-2]))
# diff_square = square1 + square2 - 2 * cross_term
# return tf.exp(-diff_square / (2 * kernel_width ** 2))
return tf.exp(-tf.reduce_sum(tf.square(x1 - x2), axis=-1) /
(2 * tf.square(kernel_width)))

def gram(self, x1, x2, kernel_width):
# x1: [..., n1, x_dim]
# x2: [..., n2, x_dim]
# kernel_width: [...]
# return: [..., n1, n2]
x_row = tf.expand_dims(x1, -2)
x_col = tf.expand_dims(x2, -3)
kernel_width = kernel_width[..., None, None]
return self.rbf_kernel(x_row, x_col, kernel_width)

def grad_gram(self, x1, x2, kernel_width):
# x1: [..., n1, x_dim]
# x2: [..., n2, x_dim]
# kernel_width: [...]
# return gram, grad_x1, grad_x2:
# [..., n1, n2], [..., n1, n2, x_dim], [..., n1, n2, x_dim]
x_row = tf.expand_dims(x1, -2)
x_col = tf.expand_dims(x2, -3)
kernel_width = kernel_width[..., None, None]
# G: [..., n1, n2]
G = self.rbf_kernel(x_row, x_col, kernel_width)
# diff: [..., n1, n2, n_x]
diff = (x_row - x_col) / (kernel_width[..., None] ** 2)
# G_expand: [..., n1, n2, 1]
G_expand = tf.expand_dims(G, axis=-1)
# grad_x1: [..., n1, n2, n_x]
grad_x2 = G_expand * diff
# grad_x2: [..., n1, n2, n_x]
grad_x1 = G_expand * (-diff)
return G, grad_x1, grad_x2

def heuristic_kernel_width(self, x_samples, x_basis):
# x_samples: [..., n_samples, x_dim]
# x_basis: [..., n_basis, x_dim]
# return: [...]
n_samples = tf.shape(x_samples)[-2]
n_basis = tf.shape(x_basis)[-2]
x_samples_expand = tf.expand_dims(x_samples, -2)
x_basis_expand = tf.expand_dims(x_basis, -3)
pairwise_dist = tf.sqrt(
tf.reduce_sum(tf.square(x_samples_expand - x_basis_expand),
axis=-1))
k = n_samples * n_basis // 2
top_k_values = tf.nn.top_k(
tf.reshape(pairwise_dist, [-1, n_samples * n_basis]),
k=k).values
kernel_width = tf.reshape(top_k_values[:, -1],
tf.shape(x_samples)[:-2])
# kernel_width = tf.Print(kernel_width, [kernel_width],
# message="kernel_width: ")
return tf.stop_gradient(kernel_width)

def compute_gradients(self, samples, x=None):
raise NotImplementedError()
@@ -0,0 +1,52 @@
#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf


def entropy_gradients(optimizer, estimator, samples, var_list=None):
dlog_q = estimator.compute_gradients(samples)
# backprop_loss = -dlog_q / tf.cast(
# tf.reduce_prod(tf.shape(samples)[:-1]), dlog_q.dtype)
# grads_and_vars = optimizer.compute_gradients(
# samples, var_list=var_list, grad_loss=backprop_loss)
surrogate_cost = tf.reduce_mean(
tf.reduce_sum(tf.stop_gradient(-dlog_q) * samples, -1))
grads_and_vars = optimizer.compute_gradients(
surrogate_cost, var_list=var_list)
return grads_and_vars


def entropy_surrogate(estimator, samples):
# surrogate of entropy - E_p \log p(x)

# sampler: A Tensor of shape [..., M, samples]
# shape(samples)
dlog_q = estimator.compute_gradients(samples)
surrogate = tf.reduce_mean(
tf.reduce_sum(tf.stop_gradient(-dlog_q) * samples, -1))
return surrogate


def minimize_entropy(optimizer, estimator, samples, var_list=None):
"""
The distribution must be reparameterizable.
The entropy will average over all dimensions before the last two dimensions.
:param optimizer: A Tensorflow Optimizer.
:param estimator: A ScoreEstimator.
:param samples: A Tensor of shape [..., M, samples]
:param var_list: A list of Variables.
:return: A Tensorflow Operation.
"""
dlog_q = estimator.compute_gradients(samples)
backprop_loss = -dlog_q / tf.cast(
tf.reduce_prod(tf.shape(samples)[:-1]), dlog_q.dtype)
opt_op = optimizer.minimize(
samples, var_list=var_list, grad_loss=backprop_loss)
return opt_op
@@ -0,0 +1,90 @@
#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf

from .base import ScoreEstimator


class SpectralScoreEstimator(ScoreEstimator):
def __init__(self, n_eigen=None, eta=None, n_eigen_threshold=None):
self._n_eigen = n_eigen
self._eta = eta
self._n_eigen_threshold = n_eigen_threshold
super(SpectralScoreEstimator, self).__init__()

def nystrom_ext(self, samples, x, eigen_vectors, eigen_values, kernel_width):
# samples: [..., M, x_dim]
# x: [..., N, x_dim]
# eigen_vectors: [..., M, n_eigen]
# eigen_values: [..., n_eigen]
# return: [..., N, n_eigen], by default n_eigen=M.
M = tf.shape(samples)[-2]
# Kxq: [..., N, M]
# grad_Kx: [..., N, M, x_dim]
# grad_Kq: [..., N, M, x_dim]
Kxq = self.gram(x, samples, kernel_width)
# Kxq = tf.Print(Kxq, [tf.shape(Kxq)], message="Kxq:")
# ret: [..., N, n_eigen]
ret = tf.sqrt(tf.to_float(M)) * tf.matmul(Kxq, eigen_vectors)
ret *= 1. / tf.expand_dims(eigen_values, axis=-2)
return ret

def compute_gradients(self, samples, x=None):
# samples: [..., M, x_dim]
# x: [..., N, x_dim]
if x is None:
kernel_width = self.heuristic_kernel_width(samples, samples)
# TODO: Simplify computation
x = samples
else:
# _samples: [..., N + M, x_dim]
_samples = tf.concat([samples, x], axis=-2)
kernel_width = self.heuristic_kernel_width(_samples, _samples)

M = tf.shape(samples)[-2]
# Kq: [..., M, M]
# grad_K1: [..., M, M, x_dim]
# grad_K2: [..., M, M, x_dim]
Kq, grad_K1, grad_K2 = self.grad_gram(samples, samples, kernel_width)
if self._eta is not None:
Kq += self._eta * tf.eye(M)
# eigen_vectors: [..., M, M]
# eigen_values: [..., M]
with tf.device("/cpu:0"):
eigen_values, eigen_vectors = tf.self_adjoint_eig(Kq)
# eigen_vectors = tf.matrix_inverse(Kq)
# eigen_values = tf.reduce_sum(Kq, -1)
# eigen_values = tf.Print(eigen_values, [eigen_values],
# message="eigen_values:", summarize=20)
if (self._n_eigen is None) and (self._n_eigen_threshold is not None):
eigen_arr = tf.reduce_mean(
tf.reshape(eigen_values, [-1, M]), axis=0)
eigen_arr = tf.reverse(eigen_arr, axis=[-1])
eigen_arr /= tf.reduce_sum(eigen_arr)
eigen_cum = tf.cumsum(eigen_arr, axis=-1)
self._n_eigen = tf.reduce_sum(
tf.to_int32(tf.less(eigen_cum, self._n_eigen_threshold)))
# self._n_eigen = tf.Print(self._n_eigen, [self._n_eigen],
# message="n_eigen:")
if self._n_eigen is not None:
# eigen_values: [..., n_eigen]
# eigen_vectors: [..., M, n_eigen]
eigen_values = eigen_values[..., -self._n_eigen:]
eigen_vectors = eigen_vectors[..., -self._n_eigen:]
# eigen_ext: [..., N, n_eigen]
eigen_ext = self.nystrom_ext(
samples, x, eigen_vectors, eigen_values, kernel_width)
# grad_K1_avg = [..., M, x_dim]
grad_K1_avg = tf.reduce_mean(grad_K1, axis=-3)
# beta: [..., n_eigen, x_dim]
beta = -tf.sqrt(tf.to_float(M)) * tf.matmul(
eigen_vectors, grad_K1_avg, transpose_a=True) / tf.expand_dims(
eigen_values, -1)
# grads: [..., N, x_dim]
grads = tf.matmul(eigen_ext, beta)
return grads
@@ -0,0 +1,57 @@
#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf

from .base import ScoreEstimator


class SteinScoreEstimator(ScoreEstimator):
def __init__(self, eta):
self._eta = eta
super(SteinScoreEstimator, self).__init__()

def compute_gradients(self, samples, x=None):
# samples: [..., M, x_dim]
# x: [..., 1, x_dim]
M = tf.shape(samples)[-2]
# kernel_width: [...]
kernel_width = self.heuristic_kernel_width(samples, samples)
# K: [..., M, M]
# grad_K1: [..., M, M, x_dim]
# grad_K2: [..., M, M, x_dim]
K, grad_K1, grad_K2 = self.grad_gram(samples, samples,
kernel_width)
# K_inv: [..., M, M]
Kinv = tf.matrix_inverse(K + self._eta * tf.eye(M))
# H_dh: [..., M, x_dim]
H_dh = tf.reduce_sum(grad_K2, axis=-2)
# grads: [..., M, x_dim]
grads = - tf.matmul(Kinv, H_dh)
if x is None:
return grads
else:
assert_single_x = tf.assert_equal(
tf.shape(x)[-2], 1,
message="Only support single-particle out-of-sample extension.")
with tf.control_dependencies([assert_single_x]):
# Kxx: [..., 1, 1]
Kxx = self.gram(x, x, kernel_width)
# Kxq: [..., 1, M]
Kxq = self.gram(x, samples, kernel_width)
# Kxq @ K_inv: [..., 1, M]
KxqKinv = tf.matmul(Kxq, Kinv)
# term1: [..., 1, 1]
term1 = -1. / (Kxx + self._eta -
tf.matmul(KxqKinv, Kxq, transpose_b=True))
# grad_Kqx2: [..., M, 1, x_dim]
Kqx, grad_Kqx1, grad_Kqx2 = self.grad_gram(samples, x, kernel_width)
# term2: [..., 1, x_dim]
term2 = tf.matmul(Kxq, grads) - tf.matmul(KxqKinv + 1.,
tf.squeeze(grad_Kqx2, -2))
# ret: [..., 1, x_dim]
return tf.matmul(term1, term2)
@@ -0,0 +1,33 @@
#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os

import numpy as np
import tensorflow as tf
from utils.utils import setup_logger

FLAGS = tf.flags.FLAGS
tf.flags.DEFINE_string("dir", "", """The result directory.""")


def main():
logger = setup_logger("eval_stats", __file__, FLAGS.dir,
filename="summary.log")
lbs = []
for seed in range(1234, 1234 + 10):
filename = os.path.join(FLAGS.dir, "eval.log.{}".format(seed))
with open(filename, "rb") as f:
text = f.read().decode("utf-8")
lb = float(text.strip().split("\n")[-1].split("=")[-1].strip())
logger.info(str(lb))
lbs.append(lb)

logger.info("{}+-{}".format(np.mean(lbs), np.std(lbs)))


if __name__ == "__main__":
main()
@@ -0,0 +1,9 @@
#!/bin/bash

for i in $(seq 1234 1243); do
echo $i;
echo $1;
echo $2;
CUDA_VISIBLE_DEVICES=$2 python -m vae.vae_eval --dir=$1 --seed=$i;
done

@@ -0,0 +1,2 @@
#!/usr/bin/env python
# -*- coding: utf-8 -*-

0 comments on commit 10db359

Please sign in to comment.
You can’t perform that action at this time.