def p_omega_net(observed=None, n_z=None, beta=1.0, mcmc_iterator=0, log_Z=0.0, initial_z=None):
net = spt.BayesianNet(observed=observed)
# sample z ~ p(z)
normal = spt.Normal(mean=tf.zeros([1, config.z_dim]),
logstd=tf.zeros([1, config.z_dim]))
normal = normal.batch_ndims_to_value(1)
z = net.add('z', normal, n_samples=n_z)
x_mean = G_omega(z)
f_z = net.add('f_z', spt.Normal(mean=x_mean, std=1.0), group_ndims=1)
x_mean = G_theta(x_mean)
x = net.add('x', spt.Normal(mean=x_mean, std=1.0), group_ndims=3)
return net
def __init__(self, mean, beta, D):
"""
Construct a new :class:`EnergyDistribution`.
Args:
mean: The mean of ExponentialDistribution.
beta: The beta of ExponentialDistribution.
"""
beta = spt.ops.convert_to_tensor_and_cast(beta, dtype=tf.float32)
_ = spt.Normal(mean, logstd=beta)
super(ExponentialDistribution,
self).__init__(dtype=tf.float32,
is_continuous=True,
is_reparameterized=True,
batch_shape=_.batch_shape,
batch_static_shape=_.get_batch_shape(),
value_ndims=_.value_ndims)
self._beta = beta
self._mean = mean
self._D = D
self._N = 1
for i in config.x_shape:
self._N *= i
self._log_Z = np.log(2) + self.N / 2.0 * np.log(
np.pi) + self.N * tf.log(self.beta)
bias = 0.0
for i in range(1, self.N):
bias += np.log(i)
for i in range(self.N - 2, 0, -2):
bias -= np.log(i / 2.0)
if self.N % 2 == 1:
bias -= 0.5 * np.log(np.pi)
self._log_Z += bias
def gaussian_mixture_prior(y, z_dim, n_clusters):
# derive the learnt z_mean
prior_mean = spt.model_variable('z_prior_mean',
dtype=tf.float32,
shape=[n_clusters, z_dim],
initializer=tf.random_normal_initializer())
z_mean = tf.nn.embedding_lookup(prior_mean, y)
# derive the learnt z_std
z_logstd = z_std = None
if config.p_z_given_y_std == 'one':
z_logstd = tf.zeros_like(z_mean)
else:
prior_std_or_logstd = spt.model_variable(
'z_prior_std_or_logstd',
dtype=tf.float32,
shape=[n_clusters, z_dim],
initializer=tf.zeros_initializer())
z_std_or_logstd = tf.nn.embedding_lookup(prior_std_or_logstd, y)
if config.p_z_given_y_std == 'one_plus_softplus_std':
z_std = 1. + tf.nn.softplus(z_std_or_logstd)
elif config.p_z_given_y_std == 'softplus_logstd':
z_logstd = tf.nn.softplus(z_std_or_logstd)
else:
assert (config.p_z_given_y_std == 'unbound_logstd')
z_logstd = z_std_or_logstd
return spt.Normal(mean=z_mean, std=z_std, logstd=z_logstd)
def q_net(x, observed=None, n_z=None, is_training=False, is_initializing=False):
net = spt.BayesianNet(observed=observed)
normalizer_fn = None if not config.act_norm else functools.partial(
spt.layers.act_norm,
axis=-1 if config.channels_last else -3,
initializing=is_initializing,
value_ndims=3,
)
# compute the hidden features
with arg_scope([spt.layers.resnet_conv2d_block],
kernel_size=config.kernel_size,
shortcut_kernel_size=config.shortcut_kernel_size,
activation_fn=tf.nn.leaky_relu,
normalizer_fn=normalizer_fn,
kernel_regularizer=spt.layers.l2_regularizer(config.l2_reg),
channels_last=config.channels_last):
h_x = tf.to_float(x)
h_x = spt.layers.resnet_conv2d_block(h_x, 16) # output: (16, 28, 28)
h_x = spt.layers.resnet_conv2d_block(h_x, 32, strides=2) # output: (32, 14, 14)
h_x = spt.layers.resnet_conv2d_block(h_x, 32) # output: (32, 14, 14)
h_x = spt.layers.resnet_conv2d_block(h_x, 64, strides=2) # output: (64, 7, 7)
h_x = spt.layers.resnet_conv2d_block(h_x, 64) # output: (64, 7, 7)
# sample z ~ q(z|x)
h_x = spt.ops.reshape_tail(h_x, ndims=3, shape=[-1])
z_mean = spt.layers.dense(h_x, config.z_dim, name='z_mean')
z_logstd = spt.layers.dense(h_x, config.z_dim, name='z_logstd')
z = net.add('z', spt.Normal(mean=z_mean, logstd=z_logstd), n_samples=n_z,
group_ndims=1)
return net
def q_net(x, observed=None, n_z=None, is_initializing=False):
net = spt.BayesianNet(observed=observed)
normalizer_fn = functools.partial(spt.layers.act_norm,
initializing=is_initializing)
# compute the hidden features
with arg_scope([spt.layers.dense],
activation_fn=tf.nn.leaky_relu,
normalizer_fn=normalizer_fn,
weight_norm=True,
kernel_regularizer=spt.layers.l2_regularizer(
config.l2_reg)):
h_x = tf.to_float(x)
h_x = spt.layers.dense(h_x, 500)
h_x = spt.layers.dense(h_x, 500)
# sample z ~ q(z|x)
z_mean = spt.layers.dense(h_x, config.z_dim, name='z_mean')
z_logstd = spt.layers.dense(h_x, config.z_dim, name='z_logstd')
z = net.add('z',
spt.Normal(mean=z_mean, logstd=z_logstd),
n_samples=n_z,
group_ndims=1)
return net
def p_net(observed=None, n_z=None, is_initializing=False):
net = spt.BayesianNet(observed=observed)
normalizer_fn = functools.partial(spt.layers.act_norm,
initializing=is_initializing)
# sample z ~ p(z)
z = net.add('z',
spt.Normal(mean=tf.zeros([1, config.z_dim]),
logstd=tf.zeros([1, config.z_dim])),
group_ndims=1,
n_samples=n_z)
# compute the hidden features
with arg_scope([spt.layers.dense],
activation_fn=tf.nn.leaky_relu,
normalizer_fn=normalizer_fn,
weight_norm=True,
kernel_regularizer=spt.layers.l2_regularizer(
config.l2_reg)):
h_z = z
h_z = spt.layers.dense(h_z, 500)
h_z = spt.layers.dense(h_z, 500)
# sample x ~ p(x|z)
x_logits = spt.layers.dense(h_z, config.x_dim, name='x_logits')
x = net.add('x', spt.Bernoulli(logits=x_logits), group_ndims=1)
return net
def q_net(x, posterior_flow, observed=None, n_z=None):
net = spt.BayesianNet(observed=observed)
normalizer_fn = None
# compute the hidden features
with arg_scope(
[spt.layers.resnet_conv2d_block],
kernel_size=config.kernel_size,
shortcut_kernel_size=config.shortcut_kernel_size,
activation_fn=tf.nn.leaky_relu,
normalizer_fn=normalizer_fn,
kernel_regularizer=spt.layers.l2_regularizer(config.l2_reg),
):
h_x = tf.to_float(x)
h_x = spt.layers.resnet_conv2d_block(
h_x, 96, kernel_size=5, scope='level_0') # output: (28, 28, 16)
h_x = spt.layers.resnet_conv2d_block(
h_x, 96, scope='level_1') # output: (14, 14, 32)
h_x = spt.layers.resnet_conv2d_block(
h_x, 192, strides=2, scope='level_2') # output: (14, 14, 32)
h_x = spt.layers.resnet_conv2d_block(
h_x, 192, scope='level_4') # output: (14, 14, 32)
h_x = spt.layers.resnet_conv2d_block(
h_x, 192, strides=2, scope='level_5') # output: (14, 14, 32)
h_x = spt.layers.resnet_conv2d_block(
h_x, 192, scope='level_7') # output: (7, 7, 64)
h_x = spt.layers.resnet_conv2d_block(
h_x, 192, scope='level_9') # output: (7, 7, 64)
z_dim_channel = 16 * config.z_dim // config.x_shape[
0] // config.x_shape[0]
z_mean = spt.layers.resnet_conv2d_block(
h_x,
z_dim_channel,
scope='z_mean',
kernel_initializer=tf.zeros_initializer())
z_logstd = spt.layers.resnet_conv2d_block(
h_x,
z_dim_channel,
scope='z_logstd',
kernel_initializer=tf.zeros_initializer())
z_mean = spt.ops.reshape_tail(z_mean, ndims=3, shape=[-1])
z_logstd = spt.ops.reshape_tail(z_logstd, ndims=3, shape=[-1])
# sample z ~ q(z|x)
# z_distribution = spt.FlowDistribution(
# spt.Normal(mean=z_mean, logstd=spt.ops.maybe_clip_value(z_logstd, min_val=config.epsilon)),
# posterior_flow
# )
# z = net.add('z', z_distribution, n_samples=n_z)
z = net.add('z',
spt.Normal(mean=z_mean,
logstd=spt.ops.maybe_clip_value(
z_logstd, min_val=config.min_logstd_of_q)),
n_samples=n_z,
group_ndims=1)
return net
def p_omega_net(observed=None,
n_z=None,
beta=1.0,
mcmc_iterator=0,
log_Z=0.0,
initial_z=None):
net = spt.BayesianNet(observed=observed)
# sample z ~ p(z)
normal = spt.Normal(mean=tf.zeros([1, 128]), logstd=tf.zeros([1, 128]))
normal = normal.batch_ndims_to_value(1)
z = net.add('z', normal, n_samples=n_z)
z_channel = 16 * config.z_dim // config.x_shape[0] // config.x_shape[1]
x_mean = G_omega(z, z_channel)
x_mean = spt.ops.reshape_tail(x_mean, ndims=3, shape=(-1, ))
f_z = net.add('f_z', spt.Normal(mean=x_mean, std=1.0), group_ndims=1)
x_mean = G_theta(x_mean)
x = net.add('x', spt.Normal(mean=x_mean, std=1.0), group_ndims=3)
return net
def make_component(i):
normal = spt.Normal(mean=tf.get_variable('mean_{}'.format(i),
shape=[1, config.z_dim],
dtype=tf.float32,
trainable=True),
logstd=tf.maximum(
tf.get_variable('logstd_{}'.format(i),
shape=[1, config.z_dim],
dtype=tf.float32,
trainable=True), -1.))
return normal.expand_value_ndims(1)
def p_net(observed=None, n_z=None, is_training=False, is_initializing=False):
net = spt.BayesianNet(observed=observed)
normalizer_fn = None if not config.act_norm else functools.partial(
spt.layers.act_norm,
axis=-1 if config.channels_last else -3,
initializing=is_initializing,
value_ndims=3,
)
# sample z ~ p(z)
z = net.add('z',
spt.Normal(mean=tf.zeros([1, config.z_dim]),
std=tf.ones([1, config.z_dim]) *
config.truncated_sigma),
group_ndims=1,
n_samples=n_z)
# compute the hidden features
with arg_scope([spt.layers.resnet_deconv2d_block],
kernel_size=config.kernel_size,
shortcut_kernel_size=config.shortcut_kernel_size,
activation_fn=tf.nn.leaky_relu,
normalizer_fn=normalizer_fn,
kernel_regularizer=spt.layers.l2_regularizer(config.l2_reg),
channels_last=config.channels_last):
h_z = spt.layers.dense(z, 64 * 7 * 7)
h_z = spt.ops.reshape_tail(h_z,
ndims=1,
shape=(7, 7,
64) if config.channels_last else
(64, 7, 7))
h_z = spt.layers.resnet_deconv2d_block(h_z, 64) # output: (64, 7, 7)
h_z = spt.layers.resnet_deconv2d_block(
h_z, 32, strides=2) # output: (32, 14, 14)
h_z = spt.layers.resnet_deconv2d_block(h_z, 32) # output: (32, 14, 14)
h_z = spt.layers.resnet_deconv2d_block(
h_z, 16, strides=2) # output: (16, 28, 28)
# sample x ~ p(x|z)
x_logits = spt.layers.conv2d(
h_z,
1, (1, 1),
padding='same',
name='feature_map_to_pixel',
channels_last=config.channels_last) # output: (1, 28, 28)
x = net.add('x',
spt.Bernoulli(logits=x_logits, dtype=tf.float32),
group_ndims=3)
return net
def p_net(observed=None, n_z=None):
net = spt.BayesianNet(observed=observed)
# sample z ~ p(z)
z = net.add('z', spt.Normal(mean=tf.zeros([1, config.z_dim]),
logstd=tf.zeros([1, config.z_dim])),
group_ndims=1, n_samples=n_z)
# compute the hidden features
with arg_scope([spt.layers.dense],
activation_fn=tf.nn.leaky_relu,
kernel_regularizer=spt.layers.l2_regularizer(config.l2_reg)):
h_z = z
h_z = spt.layers.dense(h_z, 500)
h_z = spt.layers.dense(h_z, 500)
# sample x ~ p(x|z)
x_mean = spt.layers.dense(h_z, config.x_dim, name='x_mean')
x_std = 1e-4 + tf.nn.softplus(
spt.layers.dense(h_z, config.x_dim, name='x_std'))
x = net.add('x', spt.Normal(mean=x_mean, std=x_std), group_ndims=1)
return net
def q_net(x, observed=None, n_z=None):
net = spt.BayesianNet(observed=observed)
normalizer_fn = None
# compute the hidden features
with arg_scope([spt.layers.resnet_conv2d_block],
kernel_size=config.kernel_size,
shortcut_kernel_size=config.shortcut_kernel_size,
activation_fn=tf.nn.leaky_relu,
normalizer_fn=normalizer_fn,
kernel_regularizer=spt.layers.l2_regularizer(
config.l2_reg)):
h_x = tf.to_float(x)
h_x = spt.layers.resnet_conv2d_block(
h_x, 16, scope='level_0') # output: (28, 28, 16)
h_x = spt.layers.resnet_conv2d_block(
h_x, 16, scope='level_1') # output: (28, 28, 16)
h_x = spt.layers.resnet_conv2d_block(
h_x, 32, scope='level_2') # output: (14, 14, 32)
h_x = spt.layers.resnet_conv2d_block(
h_x, 32, scope='level_3') # output: (14, 14, 32)
h_x = spt.layers.resnet_conv2d_block(
h_x, 64, strides=2, scope='level_4') # output: (14, 14, 32)
h_x = spt.layers.resnet_conv2d_block(
h_x, 64, scope='level_5') # output: (14, 14, 32)
h_x = spt.layers.resnet_conv2d_block(
h_x, 128, strides=2, scope='level_6') # output: (7, 7, 64)
h_x = spt.layers.resnet_conv2d_block(
h_x, 128, scope='level_7') # output: (7, 7, 64)
h_x = spt.layers.resnet_conv2d_block(
h_x, 256, strides=2, scope='level_8') # output: (7, 7, 64)
# sample z ~ q(z|x)
h_x = spt.ops.reshape_tail(h_x, ndims=3, shape=[-1])
z_mean = spt.layers.dense(h_x,
config.z_dim,
scope='z_mean',
kernel_initializer=tf.zeros_initializer())
z_logstd = spt.layers.dense(h_x,
config.z_dim,
scope='z_logstd',
kernel_initializer=tf.zeros_initializer())
z = net.add('z',
spt.Normal(mean=z_mean,
logstd=spt.ops.maybe_clip_value(
z_logstd, min_val=config.epsilon)),
n_samples=n_z,
group_ndims=1)
return net
def q_net(x, observed=None, n_samples=None):
net = spt.BayesianNet(observed=observed)
# compute the hidden features
with arg_scope([spt.layers.dense],
activation_fn=tf.nn.leaky_relu,
kernel_regularizer=spt.layers.l2_regularizer(
config.l2_reg)):
h_x = tf.to_float(x)
h_x = spt.layers.dense(h_x, 500)
h_x = spt.layers.dense(h_x, 500)
# sample y ~ q(y|x)
y_logits = spt.layers.dense(h_x, config.n_clusters, name='y_logits')
y = net.add('y', spt.Categorical(y_logits), n_samples=n_samples)
y_one_hot = tf.one_hot(y, config.n_clusters, dtype=tf.float32)
# sample z ~ q(z|y,x)
with arg_scope([spt.layers.dense],
activation_fn=tf.nn.leaky_relu,
kernel_regularizer=spt.layers.l2_regularizer(
config.l2_reg)):
if config.mean_field_assumption_for_q:
# by mean-field-assumption we let q(z|y,x) = q(z|x)
h_z = h_x
z_n_samples = n_samples
else:
if n_samples is not None:
h_z = tf.concat([
tf.tile(tf.reshape(h_x, [1, -1, 500]),
tf.stack([n_samples, 1, 1])), y_one_hot
],
axis=-1)
else:
h_z = tf.concat([h_x, y_one_hot], axis=-1)
h_z = spt.layers.dense(h_z, 500)
z_n_samples = None
z_mean = spt.layers.dense(h_z, config.z_dim, name='z_mean')
z_logstd = spt.layers.dense(h_z, config.z_dim, name='z_logstd')
z = net.add('z',
spt.Normal(mean=z_mean,
logstd=z_logstd,
is_reparameterized=False),
n_samples=z_n_samples,
group_ndims=1)
return net
def q_net(x, posterior_flow, observed=None, n_z=None):
net = spt.BayesianNet(observed=observed)
normalizer_fn = None
# compute the hidden features
with arg_scope([spt.layers.resnet_conv2d_block],
kernel_size=config.kernel_size,
shortcut_kernel_size=config.shortcut_kernel_size,
activation_fn=tf.nn.leaky_relu,
normalizer_fn=normalizer_fn,
kernel_regularizer=spt.layers.l2_regularizer(config.l2_reg),
dropout_fn=dropout):
h_x = tf.to_float(x)
h_x = spt.layers.resnet_conv2d_block(h_x, 16, scope='level_0') # output: (28, 28, 16)
h_x = tf.concat([h_x, x], axis=-1)
h_x = spt.layers.resnet_conv2d_block(h_x, 32, scope='level_2') # output: (14, 14, 32)
h_x = tf.concat([h_x, x], axis=-1)
h_x = spt.layers.resnet_conv2d_block(h_x, 64, scope='level_3') # output: (14, 14, 32)
h_x = tf.concat([h_x, x], axis=-1)
h_x = spt.layers.resnet_conv2d_block(h_x, 128, strides=2, scope='level_4') # output: (14, 14, 32)
x = spt.ops.reshape_tail(x, ndims=3,
shape=[config.x_shape[0] // 2, config.x_shape[1] // 2, config.x_shape[2] * 4])
h_x = tf.concat([h_x, x], axis=-1)
h_x = spt.layers.resnet_conv2d_block(h_x, 256, strides=2, scope='level_6') # output: (7, 7, 64)
x = spt.ops.reshape_tail(x, ndims=3,
shape=[config.x_shape[0] // 4, config.x_shape[1] // 4, config.x_shape[2] * 16])
h_x = tf.concat([h_x, x], axis=-1)
h_x = spt.layers.resnet_conv2d_block(h_x, 512, strides=2, scope='level_8') # output: (7, 7, 64)
x = spt.ops.reshape_tail(x, ndims=3,
shape=[config.x_shape[0] // 8, config.x_shape[1] // 8, config.x_shape[2] * 64])
h_x = tf.concat([h_x, x], axis=-1)
# sample z ~ q(z|x)
h_x = spt.ops.reshape_tail(h_x, ndims=3, shape=[-1])
# x = spt.ops.reshape_tail(x, ndims=3, shape=[-1])
# h_x = tf.concat([h_x, x], axis=-1)
z_mean = spt.layers.dense(h_x, config.z_dim, scope='z_mean', kernel_initializer=tf.zeros_initializer())
z_logstd = spt.layers.dense(h_x, config.z_dim, scope='z_logstd', kernel_initializer=tf.zeros_initializer())
z_distribution = spt.FlowDistribution(
spt.Normal(mean=z_mean, logstd=spt.ops.maybe_clip_value(z_logstd, min_val=config.epsilon)),
posterior_flow
)
z = net.add('z', z_distribution, n_samples=n_z)
# z = net.add('z', spt.Normal(mean=z_mean, logstd=spt.ops.maybe_clip_value(z_logstd, min_val=config.epsilon)),
# n_samples=n_z, group_ndims=1)
return net
def q_net(x, posterior_flow, observed=None, n_z=None):
net = spt.BayesianNet(observed=observed)
normalizer_fn = None
# compute the hidden features
with arg_scope([spt.layers.resnet_conv2d_block],
kernel_size=config.kernel_size,
shortcut_kernel_size=config.shortcut_kernel_size,
activation_fn=tf.nn.leaky_relu,
normalizer_fn=normalizer_fn,
kernel_regularizer=spt.layers.l2_regularizer(config.l2_reg), ):
h_x = tf.to_float(x)
h_x = spt.layers.resnet_conv2d_block(h_x, 16, scope='level_0') # output: (28, 28, 16)
h_x = spt.layers.resnet_conv2d_block(h_x, 32, scope='level_2') # output: (14, 14, 32)
h_x = spt.layers.resnet_conv2d_block(h_x, 32, scope='level_3') # output: (14, 14, 32)
h_x = spt.layers.resnet_conv2d_block(h_x, 32, strides=2, scope='level_4') # output: (14, 14, 32)
h_x = spt.layers.resnet_conv2d_block(h_x, 64, strides=2, scope='level_6') # output: (7, 7, 64)
# h_x = spt.layers.resnet_conv2d_block(h_x, 512, strides=2, scope='level_8') # output: (7, 7, 64)
# sample z ~ q(z|x)
h_x = spt.ops.reshape_tail(h_x, ndims=3, shape=[-1])
# x = spt.ops.reshape_tail(x, ndims=3, shape=[-1])
# h_x = tf.concat([h_x, x], axis=-1)
z_mean = spt.layers.dense(h_x, config.z_dim, scope='z_mean', kernel_initializer=tf.zeros_initializer())
z_logstd = spt.layers.dense(h_x, config.z_dim, scope='z_logstd', kernel_initializer=tf.zeros_initializer())
# sample z ~ q(z|x)
if config.use_truncated:
z_distribution = TruncatedNormal(mean=z_mean,
logstd=spt.ops.maybe_clip_value(z_logstd, min_val=config.min_logstd_of_q))
print(z_distribution.batch_shape)
else:
z_distribution = spt.Normal(mean=z_mean,
logstd=spt.ops.maybe_clip_value(z_logstd, min_val=config.min_logstd_of_q))
if config.use_flow:
z_distribution = spt.FlowDistribution(
z_distribution,
posterior_flow
)
z = net.add('z', z_distribution, n_samples=n_z)
else:
z = net.add('z', z_distribution, n_samples=n_z, group_ndims=1)
return net
def q_net(x, posterior_flow, observed=None, n_z=None):
net = spt.BayesianNet(observed=observed)
# compute the hidden features
with arg_scope([spt.layers.dense],
activation_fn=tf.nn.leaky_relu,
kernel_regularizer=spt.layers.l2_regularizer(config.l2_reg)):
h_x = tf.to_float(x)
h_x = spt.layers.dense(h_x, 500)
h_x = spt.layers.dense(h_x, 500)
# sample z ~ q(z|x)
z_mean = spt.layers.dense(h_x, config.z_dim, name='z_mean')
z_std = 1e-4 + tf.nn.softplus(
spt.layers.dense(h_x, config.z_dim, name='z_std'))
z = net.add('z', spt.Normal(mean=z_mean, std=z_std), n_samples=n_z,
group_ndims=1, flow=posterior_flow)
return net
def p_omega_net(observed=None, n_z=None, beta=1.0, mcmc_iterator=0, log_Z=0.0, initial_z=None):
net = spt.BayesianNet(observed=observed)
# sample z ~ p(z)
normal = spt.Normal(mean=tf.zeros([1, 256]),
logstd=tf.zeros([1, 256]))
normal = normal.batch_ndims_to_value(1)
xi = tf.get_variable(name='xi', shape=(), initializer=tf.constant_initializer(config.initial_xi),
dtype=tf.float32, trainable=True)
# xi = tf.square(xi)
xi = tf.nn.sigmoid(xi) # TODO
pz = EnergyDistribution(normal, G=G_omega, D=D_psi, log_Z=log_Z, xi=xi, mcmc_iterator=mcmc_iterator,
initial_z=initial_z)
z = net.add('z', pz, n_samples=n_z)
x_mean = G_omega(z)
x = net.add('x', ExponentialDistribution(
mean=x_mean,
beta=beta,
D=D_psi
), group_ndims=3)
return net
def q_net(x, posterior_flow, observed=None, n_z=None):
net = spt.BayesianNet(observed=observed)
# compute the hidden features
with arg_scope([spt.layers.dense],
activation_fn=tf.nn.leaky_relu,
kernel_regularizer=spt.layers.l2_regularizer(
config.l2_reg)):
h_x = tf.to_float(x)
h_x = spt.layers.dense(h_x, 500)
h_x = spt.layers.dense(h_x, 500)
# sample z ~ q(z|x)
z_mean = spt.layers.dense(h_x, config.z_dim, name='z_mean')
z_logstd = spt.layers.dense(h_x, config.z_dim, name='z_logstd')
z_distribution = spt.FlowDistribution(
spt.Normal(mean=z_mean, logstd=z_logstd), posterior_flow)
z = net.add('z', z_distribution, n_samples=n_z)
return net
def p_net(observed=None, n_z=None, beta=1.0, mcmc_iterator=0):
normalizer_fn = None
net = spt.BayesianNet(observed=observed)
# sample z ~ p(z)
normal = spt.Normal(mean=tf.zeros([1, config.z_dim]),
std=tf.ones([1, config.z_dim]) *
config.truncated_sigma)
normal = normal.batch_ndims_to_value(1)
pz = EnergyDistribution(normal,
G=G_theta,
D=D_psi,
log_Z=get_log_Z(),
mcmc_iterator=mcmc_iterator)
z = net.add('z', pz, n_samples=n_z)
x_mean = G_theta(z)
x = net.add('x',
ExponentialDistribution(mean=x_mean, beta=beta, D=D_psi),
group_ndims=3)
# compute the hidden features
return net
def p_net(observed=None,
n_z=None,
beta=1.0,
mcmc_iterator=0,
log_Z=0.0,
initial_z=None):
net = spt.BayesianNet(observed=observed)
# sample z ~ p(z)
normal = spt.Normal(mean=tf.zeros([1, config.z_dim]),
logstd=tf.zeros([1, config.z_dim]))
normal = normal.batch_ndims_to_value(1)
xi = tf.get_variable(name='xi',
shape=(),
initializer=tf.constant_initializer(
config.initial_xi),
dtype=tf.float32,
trainable=True)
# xi = tf.square(xi)
xi = tf.nn.sigmoid(xi) # TODO
pz = EnergyDistribution(normal,
G=G_theta,
D=D_psi,
log_Z=log_Z,
xi=xi,
mcmc_iterator=mcmc_iterator,
initial_z=initial_z)
z = net.add('z', pz, n_samples=n_z)
x_mean, x_logstd = G_theta(z, return_std=True)
x = net.add('x',
DiscretizedLogistic(
mean=x_mean,
log_scale=spt.ops.maybe_clip_value(
beta if config.uniform_scale else x_logstd,
min_val=config.epsilon),
bin_size=2.0 / 256.0,
min_val=-1.0 + 1.0 / 256.0,
max_val=1.0 - 1.0 / 256.0,
epsilon=1e-10),
n_samples=config.test_x_samples,
group_ndims=3)
return net
def p_net(observed=None, n_z=None, is_training=True, channels_last=True):
net = spt.BayesianNet(observed=observed)
# sample z ~ p(z)
z = net.add('z',
spt.Normal(mean=tf.zeros([1, config.z_dim]),
logstd=tf.zeros([1, config.z_dim])),
group_ndims=1,
n_samples=n_z)
# compute the hidden features
with arg_scope([spt.layers.resnet_deconv2d_block],
kernel_size=config.kernel_size,
shortcut_kernel_size=config.shortcut_kernel_size,
activation_fn=tf.nn.leaky_relu,
kernel_regularizer=spt.layers.l2_regularizer(config.l2_reg),
channels_last=channels_last):
h_z = spt.layers.dense(z, 64 * 7 * 7)
h_z = spt.utils.reshape_tail(
h_z, ndims=1, shape=[7, 7, 64] if channels_last else [64, 7, 7])
h_z = spt.layers.resnet_deconv2d_block(h_z, 64) # output: (64, 7, 7)
h_z = spt.layers.resnet_deconv2d_block(
h_z, 32, strides=2) # output: (32, 14, 14)
h_z = spt.layers.resnet_deconv2d_block(h_z, 32) # output: (32, 14, 14)
h_z = spt.layers.resnet_deconv2d_block(
h_z, 16, strides=2) # output: (16, 28, 28)
# sample x ~ p(x|z)
h_z = spt.layers.conv2d(h_z,
1, (1, 1),
padding='same',
name='feature_map_to_pixel',
channels_last=channels_last) # output: (1, 28, 28)
x_logits = spt.utils.reshape_tail(h_z, 3, [config.x_dim])
x = net.add('x', spt.Bernoulli(logits=x_logits), group_ndims=1)
return net
def q_net(x, observed=None, n_z=None, is_initializing=False):
net = spt.BayesianNet(observed=observed)
normalizer_fn = functools.partial(spt.layers.act_norm,
initializing=is_initializing)
# compute the hidden features
with arg_scope([spt.layers.dense],
activation_fn=tf.nn.leaky_relu,
normalizer_fn=normalizer_fn,
weight_norm=True,
kernel_regularizer=spt.layers.l2_regularizer(
config.l2_reg)):
h_x = tf.to_float(x)
h_x = spt.layers.dense(h_x, 500)
h_x = spt.layers.dense(h_x, 500)
# sample z ~ q(z|x)
components = [
spt.Normal(mean=spt.layers.dense(h_x,
config.z_dim,
name='z_mean_{}'.format(i)),
logstd=spt.layers.dense(h_x,
config.z_dim,
name='z_logstd_{}'.format(i)))
for i in range(config.n_mixture_components)
]
mixture_param_shape = spt.utils.concat_shapes([
spt.utils.get_shape(components[0].mean), [config.n_mixture_components]
])
mixture = spt.Mixture(
categorical=spt.Categorical(logits=tf.zeros(mixture_param_shape)),
components=components,
is_reparameterized=True)
z = net.add('z', mixture, n_samples=n_z, group_ndims=1)
return net
def p_net(observed=None, n_z=None):
net = spt.BayesianNet(observed=observed)
# sample z ~ p(z)
z = net.add('z',
spt.Normal(mean=tf.zeros([1, config.z_dim]),
logstd=tf.zeros([1, config.z_dim])),
group_ndims=1,
n_samples=n_z)
# compute the hidden features
with arg_scope([spt.layers.dense],
activation_fn=tf.nn.leaky_relu,
kernel_regularizer=spt.layers.l2_regularizer(
config.l2_reg)):
h_z = z
h_z = spt.layers.dense(h_z, 500, scope='hidden_0')
h_z = spt.layers.dense(h_z, 500, scope='hidden_1')
# sample x ~ p(x|z)
x_logits = spt.layers.dense(h_z, config.x_dim, scope='x_logits')
x = net.add('x', spt.Bernoulli(logits=x_logits), group_ndims=1)
return net
def p_net(observed=None, n_z=None, beta=1.0, mcmc_iterator=0, log_Z=0.0, initial_z=None,
mcmc_alpha=config.smallest_step):
net = spt.BayesianNet(observed=observed)
# sample z ~ p(z)
normal = spt.Normal(mean=tf.zeros([1, config.z_dim]),
logstd=tf.zeros([1, config.z_dim]))
normal = normal.batch_ndims_to_value(1)
xi = tf.get_variable(name='xi', shape=(), initializer=tf.constant_initializer(config.initial_xi),
dtype=tf.float32, trainable=True)
# xi = tf.square(xi)
xi = tf.nn.sigmoid(xi) # TODO
pz = EnergyDistribution(normal, G=G_theta, D=D_psi, log_Z=log_Z, xi=xi, mcmc_iterator=mcmc_iterator,
initial_z=initial_z, mcmc_alpha=mcmc_alpha)
z = net.add('z', pz, n_samples=n_z)
x_mean = G_theta(z)
x_mean = tf.clip_by_value(x_mean, 1e-7, 1 - 1e-7)
logits = tf.log(x_mean) - tf.log1p(-x_mean)
bernouli = spt.Bernoulli(
logits=logits, dtype=tf.float32
)
# bernouli.mean = x_mean
x = net.add('x', bernouli, group_ndims=3)
return net
def q_net(x,
observed=None,
n_z=None,
is_training=False,
is_initializing=False):
"""
Inference net
param x: input X, multivariate time series data.
return q net structure.
"""
net = spt.BayesianNet(observed=observed)
normalizer_fn = None if not config.act_norm else functools.partial(
spt.layers.act_norm,
axis=-1 if config.channels_last else -3,
initializing=is_initializing,
value_ndims=3,
)
print("=" * 10 + "qnet" + "=" * 10)
# compute the hidden features
with arg_scope([spt.layers.resnet_conv2d_block],
kernel_size=config.kernel_size2,
shortcut_kernel_size=config.shortcut_kernel_size,
activation_fn=tf.nn.elu,
normalizer_fn=normalizer_fn,
kernel_regularizer=spt.layers.l2_regularizer(config.l2_reg),
channels_last=config.channels_last):
print("qx:%s" % x.get_shape())
h_x = tf.reshape(tf.to_float(x),
[-1, config.timeLength, config.metricNumber, 1]
if config.channels_last else
[-1, 1, config.timeLength, config.metricNumber])
print("q1:%s" % h_x.get_shape())
h_x = spt.layers.resnet_conv2d_block(h_x,
1,
kernel_size=(config.kernel_size1,
1),
strides=(config.strides1, 1))
print("q2:%s" % h_x.get_shape())
h_x = spt.layers.resnet_conv2d_block(h_x,
1,
kernel_size=(config.kernel_size1,
1),
strides=(config.strides1, 1))
print("q3:%s" % h_x.get_shape())
h_x = spt.layers.resnet_conv2d_block(h_x,
1,
kernel_size=(config.kernel_size2,
1),
strides=(config.strides2, 1))
print("q4:%s" % h_x.get_shape())
h_x = spt.layers.resnet_conv2d_block(h_x,
1,
kernel_size=(config.kernel_size2,
1),
strides=(config.strides2, 1))
print("q5:%s" % h_x.get_shape())
h_x = spt.ops.reshape_tail(h_x, ndims=3, shape=[-1])
print("q6:%s" % h_x.get_shape())
# sample y ~ q(y|x)
c_logits = spt.layers.dense(h_x, config.n_c, name='c_logits')
c = net.add('c', spt.Categorical(c_logits))
c_one_hot = tf.one_hot(c, config.n_c, dtype=tf.float32)
print("qc:%s, %s, %s" % (c_logits.shape, c.shape, c_one_hot.shape))
h_z = h_x
# sample z ~ q(z|x)
z_mean = spt.layers.dense(h_z, config.z_dim, name='z_mean')
z_logstd = spt.layers.dense(
h_z, config.z_dim, name='z_logstd',
activation_fn=tf.nn.elu) + config.std_epsilon
z = net.add('z',
spt.Normal(mean=z_mean, logstd=z_logstd),
n_samples=n_z,
group_ndims=1)
print("q7:%s, %s, %s" %
(z_mean.get_shape(), z_logstd.get_shape(), z.get_shape()))
return net
def p_net(observed=None, n_z=None, is_training=False, is_initializing=False):
"""
Generative net
return p net structure.
"""
net = spt.BayesianNet(observed=observed)
normalizer_fn = None if not config.act_norm else functools.partial(
spt.layers.act_norm,
axis=-1 if config.channels_last else -3,
initializing=is_initializing,
value_ndims=3,
)
def make_component(i):
normal = spt.Normal(mean=tf.get_variable('mean_{}'.format(i),
shape=[1, config.z_dim],
dtype=tf.float32,
trainable=True),
logstd=tf.maximum(
tf.get_variable('logstd_{}'.format(i),
shape=[1, config.z_dim],
dtype=tf.float32,
trainable=True), -1.))
return normal.expand_value_ndims(1)
components = [make_component(i) for i in range(config.n_c)]
mixture = spt.Mixture(
categorical=spt.Categorical(logits=tf.zeros([1, config.n_c])),
components=components,
is_reparameterized=True)
z = net.add('z', mixture, n_samples=n_z)
print("=" * 10 + "pnet" + "=" * 10)
# compute the hidden features
with arg_scope([spt.layers.resnet_deconv2d_block],
kernel_size=config.kernel_size2,
shortcut_kernel_size=config.shortcut_kernel_size,
activation_fn=tf.nn.elu,
normalizer_fn=normalizer_fn,
kernel_regularizer=spt.layers.l2_regularizer(config.l2_reg),
channels_last=config.channels_last):
print("px:%s" % z.get_shape())
h_z = spt.layers.dense(
z,
int(config.timeLength / (config.strides1**2) /
(config.strides2**2) * int(config.metricNumber)))
h_z = spt.ops.reshape_tail(
h_z,
ndims=1,
shape=(int(config.timeLength /
(config.strides1**2) / (config.strides2**2)),
int(config.metricNumber), 1) if config.channels_last else
(1,
int(config.timeLength / (config.strides1**2) /
(config.strides2**2)), int(config.metricNumber)))
print("p1:%s" % h_z.get_shape())
h_z = spt.layers.resnet_deconv2d_block(
h_z,
1,
kernel_size=(config.kernel_size2, 1),
strides=(config.strides2, 1))
print("p2:%s" % h_z.get_shape())
h_z = spt.layers.resnet_deconv2d_block(
h_z,
1,
kernel_size=(config.kernel_size2, 1),
strides=(config.strides2, 1))
print("p3:%s" % h_z.get_shape())
h_z = spt.layers.resnet_deconv2d_block(
h_z,
1,
kernel_size=(config.kernel_size1, 1),
strides=(config.strides1, 1))
print("p4:%s" % h_z.get_shape())
h_z = spt.layers.resnet_deconv2d_block(
h_z,
1,
kernel_size=(config.kernel_size1, 1),
strides=(config.strides1, 1))
print("p5:%s" % h_z.get_shape())
# sample x ~ p(x|z)
x_mean = spt.layers.conv2d(h_z,
1, (1, 1),
padding='same',
name='x_mean',
channels_last=config.channels_last)
x_logstd = spt.layers.conv2d(
h_z,
1,
(1, 1),
padding='same',
name='x_logstd',
channels_last=config.channels_last,
activation_fn=tf.nn.elu,
) + config.std_epsilon
x = net.add('x',
spt.Normal(mean=x_mean, logstd=x_logstd),
n_samples=n_z,
group_ndims=3)
print("p6:%s, %s, %s" %
(x_mean.get_shape(), x_logstd.get_shape(), x.get_shape()))
return net
