def log_likelihoods(self, samples, name=None):
name = adapt_name(name, "log-likelihood_%s" % self.name)
with tf.name_scope(name):
out = np.log(2 * np.pi) + self.log_variances_with_L
out += tf.square(samples - self.means_with_L) / self.variances_with_L
out *= -0.5
return tf.identity(out, name)
def sample(self, L=1, name=None):
name = adapt_name(name, "sample_%s" % self.name)
with tf.name_scope(name):
shape = tf.concat([[self.N, L], self.dims], axis=0)
noise = tf.random_normal(shape, 0, 1, dtype=tf.float32, name="noise")
samples = self.means_with_L + (self.log_variances_with_L * noise)
return tf.identity(samples, name)
def reducechoice(inp, axis, name=None):
name = adapt_name(name, "reducechoice")
with tf.name_scope(name):
avg = tfl.layers.core.flatten(tf.reduce_mean(inp, axis), name="avg")
maxx = tfl.layers.core.flatten(tf.reduce_max(inp, axis), name="max")
minn = tfl.layers.core.flatten(tf.reduce_min(inp, axis), name="min")
pooled = choice([avg, maxx, minn], name=name)
return pooled
def split_and_recombined(inps, fn, num_splits, name=None):
name = adapt_name(name, "split-and-recombine")
with tf.name_scope(name):
adapted_inps = []
# Split inputs:
with tf.name_scope("preprocessing"):
for inp in inps:
if isinstance(inp, list) or isinstance(inp, tuple):
if len(inp) % num_splits != 0:
raise RuntimeError(
"List not divisible by number of splits: %s" %
repr(inp))
stride = len(inp) // num_splits
squeeze = lambda x: x[0] if len(x) == 1 else x
adapted_inps.append([
squeeze(inp[i:(i + stride)])
for i in range(0, len(inp), stride)
])
elif (isinstance(inp, tf.Variable) or isinstance(inp, tf.Tensor))\
and len(inp.shape) > 0:
if inp.shape[0].value is None:
raise RuntimeError(
"Batch index must be defined for tensor")
leng = int(inp.shape[0])
if leng % num_splits != 0:
raise RuntimeError(
"Tensor not divisible by number of splits (%d): %s"
% (num_splits, inp.shape))
stride = leng // num_splits
adapted_inps.append([
tf.slice(inp, [
i if j == 0 else 0 for j in range(len(inp.shape))
], [
stride if j == 0 else -1
for j in range(len(inp.shape))
]) for i in range(0, leng, stride)
])
else:
adapted_inps.append([inp] * num_splits)
# Zip inputs to divide work:
adapted_inps = list(zip(*adapted_inps))
# Do work
raw_outputs = []
for split, args in enumerate(adapted_inps):
with tf.name_scope("bin%d" % split):
raw_outputs.append(fn(*args))
# Post-process outputs
outputs = []
with tf.name_scope("postprocessing"):
for i, group in enumerate(raw_outputs):
for j, var in enumerate(group):
if i == 0:
outputs.append([var])
else:
outputs[j].append(var)
return outputs
def kl_divergence(cls, p, q, name=None):
name = adapt_name(name, "kl-divergence")
with tf.name_scope(name):
inner = p.variances + tf.square(p.means - q.means)
inner /= q.variances
inner = 1 + p.log_variances - q.log_variances - inner
inner *= -0.5
kl = tf.reduce_sum(inner, list(range(1, len(inner.shape))))
kl = tf.identity(kl, name)
return kl
def SpatialMaxPooling(inp, kW, kH, dW=1, dH=1, padW=0, padH=0, **kwargs):
name = adapt_name(kwargs.get("name", None), "pool")
with tf.name_scope(name):
out = inp
out = pad(out, padW, padH)
config = dict(
strides=(dW, dH),
padding='valid',
)
config.update(kwargs)
out = tfl.layers.conv.max_pool_2d(out, (kW, kH), **config)
return out
def merge_grads(tower_grads, name=None):
name = adapt_name(name, "merge-grads")
with tf.name_scope(name):
grads = []
for grad_group in zip(*tower_grads):
var = grad_group[0][1]
vals = [val for val, var in grad_group if val is not None]
if len(vals) == 0:
grads.append((None, var))
else:
grads.append((tf.reduce_mean(tf.stack(vals, 0), 0), var))
return grads
def __init__(self, means, log_variances, name=None):
name = adapt_name(name, "gaussian")
with tf.name_scope(name):
self.name = name
self.N = tf.identity(tf.shape(means)[0], "N")
self.dims = tf.shape(means)[1:]
self.means = tf.identity(means, "mean")
self.means_with_L = tf.expand_dims(means, 1)
self.log_variances = tf.identity(log_variances, "log-variance")
self.log_variances_with_L = tf.expand_dims(log_variances, 1)
self.stddevs = tf.identity(tf.exp(0.5 * self.log_variances), "stddev")
self.variances = tf.identity(tf.exp(self.log_variances), "variance")
self.variances_with_L = tf.expand_dims(self.variances, 1)
def SpatialConvolution(inp, _, nfilters, kW, kH, dW, dH, padW, padH, **kwargs):
name = adapt_name(kwargs.get("name", None), "conv")
with tf.variable_scope(name):
out = inp
out = pad(out, padW, padH)
config = dict(
strides=(dW, dH),
padding='valid',
regularizer='L2',
weights_init='xavier',
bias_init='zeros',
weight_decay=1.0,
)
config.update(kwargs)
out = tfl.layers.conv.conv_2d(out, nfilters, (kW, kH), **config)
return out
def choice(inps, name=None):
name = adapt_name(name, "choice")
with tf.variable_scope(name):
assert len(inps) > 0, "You must provide at least one input."
shape = [len(inps)] + [int(x) for x in inps[0].shape[1:]]
W = tf.get_variable("W",
shape=shape,
dtype=tf.float32,
initializer=tf.constant_initializer(1.0),
trainable=True)
mask = tf.expand_dims(tf.nn.softmax(W, dim=0), 0)
stacked_inps = tf.stack(inps, axis=1)
out = tf.reduce_sum(tf.multiply(stacked_inps, mask), axis=1)
out = tf.identity(out, name)
print(out)
return out
def pad(tensor, padW, padH, name=None):
name = adapt_name(name, "pad")
return tf.pad(tensor, [[0, 0], [padW, padW], [padH, padH], [0, 0]],
name=name)
def symmetric_kl_divergence(cls, p, q, name=None):
name = adapt_name(name, "symmetric-kl-divergence")
with tf.name_scope(name):
return tf.identity(cls.kl_divergence(p, q) + cls.kl_divergence(q, p), name)
def kl_divergence_from_unit(self, name=None):
name = adapt_name(name, "kl-divergence-from-unit_%s" % self.name)
return DiagonalCovarianceGaussian.kl_divergence(self,
self.unit_gaussian(), name=name)
def unit_gaussian(self, name=None):
name = adapt_name(name, "unit-gaussian")
with tf.name_scope(name):
return DiagonalCovarianceGaussian(
tf.zeros(tf.shape(self.means)),
tf.ones(tf.shape(self.log_variances)), name=name)
def splitfn(inp, fn, maxbatch=None, allow_unrolling=True, name=None):
name = adapt_name(name, "splitfn")
with tf.variable_scope(name) as scope:
if not allow_unrolling or inp.shape[0].value is None:
leng = tf.shape(inp)[0]
def minibatch():
scope.reuse_variables()
remainder = tf.mod(leng, maxbatch, name="remainder")
splits = tf.identity(tf.floor_div(leng - remainder, maxbatch),
"splits")
remainder_inp = tf.slice(inp, [
leng - remainder if i == 0 else 0
for i in range(len(inp.shape))
], [-1 for i in range(len(inp.shape))])
majority_inp = tf.slice(inp,
[0 for i in range(len(inp.shape))], [
leng - remainder if i == 0 else -1
for i in range(len(inp.shape))
])
split_inp = tf.reshape(
majority_inp,
tf.concat([[splits, maxbatch],
tf.shape(inp)[1:]], 0))
majority_out = tf.map_fn(fn, split_inp)
scope.reuse_variables()
remainder_out = fn(remainder_inp)
out = tf.concat([
tf.reshape(
majority_out,
tf.concat([[leng - remainder],
tf.shape(majority_out)[2:]], 0)),
remainder_out
], 0)
if inp.shape[0].value is not None:
out = tf.reshape(
out,
tf.concat([[int(inp.shape[0])],
tf.shape(out)[1:]], 0))
return out
if maxbatch is None:
out = fn(inp)
else:
out = tf.case([(maxbatch < leng, minibatch)], lambda: fn(inp))
else:
leng = int(inp.shape[0])
if maxbatch is not None and maxbatch < leng:
remainder = leng % maxbatch
splits = (leng - remainder) // maxbatch
remainder_inp = tf.slice(inp, [
leng - remainder if i == 0 else 0
for i in range(len(inp.shape))
], [-1 for i in range(len(inp.shape))])
majority_inp = tf.slice(inp,
[0 for i in range(len(inp.shape))], [
leng - remainder if i == 0 else -1
for i in range(len(inp.shape))
])
split_inp = tf.reshape(
majority_inp,
tf.concat([[splits, maxbatch],
tf.shape(inp)[1:]], 0))
majority_out = tf.map_fn(fn, split_inp)
scope.reuse_variables()
remainder_out = fn(remainder_inp)
out = tf.concat([
tf.reshape(
majority_out,
tf.concat([[leng - remainder],
tf.shape(majority_out)[2:]], 0)),
remainder_out
], 0)
else:
out = fn(inp)
return tf.identity(out, name)
