def __init__(self, mu, sigma, dim, static_data=None):
super(SizedMvnGaussDataset, self).__init__(dim,
static_data=static_data)
self.mu = mu
self.sigma = sigma
self.data_dist = DiagMvn(mu=[self.mu] * self.dim,
sigma=[self.sigma] * self.dim)
self.num_samples = 1000
self.data = self.data_dist.get_samples(num_samples=self.num_samples)
class MvnGaussDataset(ToyDataset):
def __init__(self, mu, sigma, dim, static_data=None):
super(MvnGaussDataset, self).__init__(dim, static_data=static_data)
self.mu = mu
self.sigma = sigma
self.data_dist = DiagMvn(mu=[self.mu] * self.dim,
sigma=[self.sigma] * self.dim)
def gen_batch(self, batch_size):
return self.data_dist.get_samples(num_samples=batch_size)
class SizedMvnGaussDataset(ToyDataset):
def __init__(self, mu, sigma, dim, static_data=None):
super(SizedMvnGaussDataset, self).__init__(dim,
static_data=static_data)
self.mu = mu
self.sigma = sigma
self.data_dist = DiagMvn(mu=[self.mu] * self.dim,
sigma=[self.sigma] * self.dim)
self.num_samples = 1000
self.data = self.data_dist.get_samples(num_samples=self.num_samples)
def gen_batch(self, batch_size):
inds = np.random.choice(np.arange(self.num_samples), size=batch_size)
return self.data[inds, :].reshape((batch_size, ) + self.data.shape[1:])
def build_graph(cmd_args, db, energy_func, flow):
ph_real_data = tf.placeholder(tf.float32, shape=(None, db.dim))
ph_z0 = tf.placeholder(tf.float32, shape=(None, db.dim))
init_dist = DiagMvn(mu=[0.0] * db.dim, sigma=[1.0] * db.dim)
gen_loss, train_gen = get_gen_loss(cmd_args, ph_real_data, ph_z0, flow,
energy_func, init_dist)
disc_loss, train_disc = get_disc_loss(cmd_args, ph_real_data, ph_z0, flow,
energy_func, init_dist)
ema = tf.train.ExponentialMovingAverage(decay=cmd_args.ema_decay)
with tf.control_dependencies([train_disc]):
vars = tf.trainable_variables(scope='energy_func')
disc_train_op = ema.apply(vars)
train_disc = disc_train_op
return ph_real_data, ph_z0, init_dist, gen_loss, train_gen, disc_loss, train_disc, ema
print('plot size:', plt_size)
# plot true data
data_gen = db.data_gen(batch_size=1000, auto_reset=False)
true_data = next(data_gen)
output_path = os.path.join(cmd_args.save_dir, 'ground_truth.pdf')
plot_samples(true_data, output_path)
gamma = get_gamma(true_data, cmd_args.mmd_bd)
eval_metric = lambda x: MMD(true_data, x, gamma)
graph_list = build_graph(cmd_args, db, energy_func, flow)
ema = graph_list[-1]
ema_energy_func = build_ema_f(cmd_args, db, ema)
# computation graph for getting samples
init_dist = DiagMvn(mu=[0.0] * db.dim, sigma=[1.0] * db.dim)
ph_z0 = tf.placeholder(tf.float32, shape=(None, db.dim))
log_pz0 = init_dist.get_log_pdf(ph_z0)
x_samples, _ = flow(ph_z0, log_pz0)
# computation graph for getting pdf
factors = [0.01, 0.05, 0.1, 0.2, 0.5, 1, 5]
pdf_dict = {}
ph_x = tf.placeholder(tf.float32, shape=(None, db.dim))
x_energy = tf.reshape(-energy_func(ph_x), shape=[-1])
for key in factors:
pdf_dict[key] = tf.nn.softmax(key * x_energy)
x_energy_ema = tf.reshape(-ema_energy_func(ph_x), shape=[-1])
pdf_ema_dict = {}
def __init__(self, mu, sigma, dim, static_data=None):
super(MvnGaussDataset, self).__init__(dim, static_data=static_data)
self.mu = mu
self.sigma = sigma
self.data_dist = DiagMvn(mu=[self.mu] * self.dim,
sigma=[self.sigma] * self.dim)
from ade.common.distributions import DiagMvn
import ade.common.mcmc_net as mcmc_net
if __name__ == '__main__':
random.seed(cmd_args.seed)
np.random.seed(cmd_args.seed)
tf.random.set_random_seed(cmd_args.seed)
true_mean = -1.0
true_std = 0.1
with tf.variable_scope('energy_func'):
energy_func = GaussianEnergy(dim=cmd_args.gauss_dim, init_mu=true_mean, init_sigma=true_std)
with tf.variable_scope('generator'):
mcmc = mcmc_net.HMCNet(cmd_args, dim=cmd_args.gauss_dim)
ph_z0 = tf.placeholder(tf.float32, shape=(1000, cmd_args.gauss_dim))
zt = mcmc(energy_func, ph_z0)
init_dist = DiagMvn(mu=[0.0] * cmd_args.gauss_dim,
sigma=[1.0] * cmd_args.gauss_dim)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
z0 = init_dist.get_samples(num_samples=1000)
z_samples = sess.run(zt, feed_dict={ph_z0: z0})
print(np.mean(z_samples, axis=0))
print(np.std(z_samples, axis=0))