def p_net(observed=None, n_y=None, n_z=None, n_samples=None): if n_samples is not None: warnings.warn('`n_samples` is deprecated, use `n_y` instead.') n_y = n_samples net = spt.BayesianNet(observed=observed) # sample y y = net.add('y', spt.Categorical(tf.zeros([1, config.n_clusters])), n_samples=n_y) # sample z ~ p(z|y) z = net.add('z', gaussian_mixture_prior(y, config.z_dim, config.n_clusters), group_ndims=1, n_samples=n_z, is_reparameterized=False) # compute the hidden features for x with arg_scope([spt.layers.dense], activation_fn=tf.nn.leaky_relu, kernel_regularizer=spt.layers.l2_regularizer( config.l2_reg)): h_x = z h_x = spt.layers.dense(h_x, 500) h_x = spt.layers.dense(h_x, 500) # sample x ~ p(x|z) x_logits = spt.layers.dense(h_x, config.x_dim, name='x_logits') x = net.add('x', spt.Bernoulli(logits=x_logits), group_ndims=1) return net
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, 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 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 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, 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_logits = spt.layers.dense(h_x, config.z_dim, name='z_logits') z = net.add('z', spt.Bernoulli(logits=z_logits), 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) 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), config.z_logstd_min)) return normal.expand_value_ndims(1) components = [ make_component(i) for i in range(config.n_mixture_components) ] mixture = spt.Mixture(categorical=spt.Categorical( logits=tf.zeros([1, config.n_mixture_components])), components=components, is_reparameterized=True) z = net.add('z', mixture, 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, 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 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 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 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_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 p_net(observed=None, n_z=None): net = spt.BayesianNet(observed=observed) # sample z ~ p(z) z = net.add('z', spt.Bernoulli(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 = tf.to_float(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 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): 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 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, 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