def write_images(H, ema_params, viz_batch_original, viz_batch_processed, fname,
logprint):
rng = random.PRNGKey(H.seed_sample)
ema_apply = partial(
VAE(H).apply, {'params': jax_utils.unreplicate(ema_params)})
forward_get_latents = partial(ema_apply, method=VAE(H).forward_get_latents)
forward_samples_set_latents = partial(
ema_apply, method=VAE(H).forward_samples_set_latents)
forward_uncond_samples = partial(ema_apply,
method=VAE(H).forward_uncond_samples)
zs = [s['z'] for s in forward_get_latents(viz_batch_processed, rng)]
batches = [viz_batch_original.numpy()]
mb = viz_batch_processed.shape[0]
lv_points = np.floor(
np.linspace(0, 1, H.num_variables_visualize + 2) *
len(zs)).astype(int)[1:-1]
for i in lv_points:
batches.append(forward_samples_set_latents(mb, zs[:i], rng, t=0.1))
for t in [1.0, 0.9, 0.8, 0.7][:H.num_temperatures_visualize]:
batches.append(forward_uncond_samples(mb, rng, t=t))
n_rows = len(batches)
im = np.concatenate(batches, axis=0).reshape(
(n_rows, mb,
*viz_batch_processed.shape[1:])).transpose([0, 2, 1, 3, 4]).reshape([
n_rows * viz_batch_processed.shape[1],
mb * viz_batch_processed.shape[2], 3
])
logprint(f'printing samples to {fname}')
Image.fromarray(im).save(fname)
def load_vaes(H, logprint):
vae = VAE(H)
if H.restore_path:
logprint(f'Restoring vae from {H.restore_path}')
restore_params(vae, H.restore_path, map_cpu=True, local_rank=H.local_rank, mpi_size=H.mpi_size)
ema_vae = VAE(H)
if H.restore_ema_path:
logprint(f'Restoring ema vae from {H.restore_ema_path}')
restore_params(ema_vae, H.restore_ema_path, map_cpu=True, local_rank=H.local_rank, mpi_size=H.mpi_size)
else:
ema_vae.load_state_dict(vae.state_dict())
ema_vae.requires_grad_(False)
vae = vae.cuda(H.local_rank)
ema_vae = ema_vae.cuda(H.local_rank)
vae = DistributedDataParallel(vae, device_ids=[H.local_rank], output_device=H.local_rank)
if len(list(vae.named_parameters())) != len(list(vae.parameters())):
raise ValueError('Some params are not named. Please name all params.')
total_params = 0
for name, p in vae.named_parameters():
total_params += np.prod(p.shape)
logprint(total_params=total_params, readable=f'{total_params:,}')
return vae, ema_vae
def encoding_process(load_file_name, file_list):
import tensorflow as tf
file_name_list = [os.path.splitext(file)[0] for file in file_list]
data_length = len(file_list)
# Load models.
vae = VAE()
vae.build_model(is_training=False, is_assigning=False)
with tf.Session(graph=vae.graph) as sess:
# Load variables.
saver = tf.train.Saver()
saver.restore(sess, SAVE_VAE_DIR + file_name)
for i in range(data_length):
# Load raw data.
data = np.load(RAW_DATA_DIR + file_list[i])
obs = data["obs"]
action = data["action"]
# Compute the mean and standard deviation rather than the encoding z.
mu, sigma = sess.run([vae.mu, vae.sigma],
feed_dict={vae.Input: obs / 255.0})
# Save file.
np.savez_compressed(ENCODING_DATA_DIR + file_name_list[i],
mu=mu,
sigma=sigma,
action=action)
def main(args):
transform = transforms.Compose([
transforms.Resize([64, 64], 1),
transforms.ToTensor(),
transforms.Normalize([.5], [.5])
])
vae = VAE(img_size=[3, 64, 64], z_dim=args.z_dim)
pt_file = load_model(args.model_load_path, "*.pt")
vae.load_state_dict(torch.load(pt_file))
vae.eval()
dirs = glob.glob(os.path.join(args.data_path, "goal*/camera*"))
dirs = [d for d in dirs if os.path.isdir(d)]
dirs.sort()
for d in dirs:
data = []
files = glob.glob(os.path.join(d, "*.png"))
files.sort()
print(d.split("/")[-1])
for f in files:
img = Image.open(f)
img = transform(img)
data.append(vae.reparam(*vae.encoder(img[None, :, :, :])).detach())
data = torch.cat(data).cpu().numpy()
print(d + ".txt")
np.savetxt(d + ".txt", data, delimiter=",")
def train(self,
generations=1500,
evolutionary_leap_gens=10,
disp_best=True):
env = gym.make('CarRacing-v0', verbose=0)
env = wrap_env(env, W, H, gray_scale=False)
latent_size = 32
vae_model = VAE(latent_size)
checkpoint_dir = './vae_ckpt/'
latest = tf.train.latest_checkpoint(checkpoint_dir)
vae_model.load_weights(latest)
for g in range(1, generations + 1):
print('Generation', g)
st = time()
if not g % evolutionary_leap_gens:
print('Leap')
best_model = self.evolutionary_leap()
else:
best_model = self.next_gen()
print(best_model)
if disp_best:
best_model.evaluate(env, vae_model, disp=True, n=1)
best_model.save_all(name='{generation:04d}'.format(generation=g) +
'result:' + str(best_model.mean_result))
print('T:', time() - st)
env.close()
def __init__(self, n_input, n_hidden, dim_z, n_output, gamma, binary=True, **kwargs):
"""initialize neural networks
:param gamma: weight for total correlation term in loss function
"""
super(FactorVAE, self).__init__()
self.dim_z = dim_z
self.binary = binary
self.gamma = gamma
self.input_size = (n_input,)
# VAE networks
self.vae = VAE(n_input, n_hidden, dim_z, n_output, binary, **kwargs)
# discriminator layers
D_hidden_num = 3
D_hidden_dim = 1000
D_hidden_dims = [D_hidden_dim] * D_hidden_num
D_act = nn.LeakyReLU
D_act_args = {"negative_slope": 0.2, "inplace": False}
D_output_dim = 2
self.discriminator = nns.create_mlp(self.dim_z, D_hidden_dims,
act_layer=D_act, act_args=D_act_args, norm=True)
self.discriminator = nn.Sequential(
self.discriminator,
nn.Linear(D_hidden_dim, D_output_dim))
def vae_visualization(file_name="vae", random_file=None):
# Create folders.
if not os.path.isdir(FIGURE_VAE_VISUALIZATION_DIR):
os.makedirs(FIGURE_VAE_VISUALIZATION_DIR)
if random_file == None:
# Load random data.
file_list = os.listdir(RAW_DATA_DIR)
random_file = np.random.choice(file_list)
random_file_name = os.path.splitext(random_file)[0]
obs = np.load(RAW_DATA_DIR + random_file)["obs"]
# Load models.
vae = VAE()
vae.build_model(is_training=False, is_assigning=False)
with tf.Session(graph=vae.graph) as sess:
# Load variables.
saver = tf.train.Saver()
saver.restore(sess, SAVE_VAE_DIR + file_name)
# Compute the reconstruction.
recons = sess.run(vae.output, feed_dict={vae.Input: obs / 255.0})
tf.contrib.keras.backend.clear_session()
imageio.mimsave(
FIGURE_VAE_VISUALIZATION_DIR + random_file_name + ".gif",
[plot_obs_recons(obs[i], recons[i]) for i in range(MAX_FRAME)],
fps=20)
def train_model_cifar(model_dir,
num_steps,
batch_size=64,
learning_rate=0.0005):
"""
Train a VAE model
:param model_dir:
:param num_steps:
:param batch_size:
:param learning_rate:
:return:
"""
with tf.Session() as sess:
vae = VAE(sess=sess,
model_dir=model_dir,
batch_size=batch_size,
learning_rate=learning_rate)
# Load cifar dataset
load_cifar_data(
cifar_path=
"/home/kevin/deep_learning/cifar-10-python/cifar-10-batches-py/",
batch_size=batch_size)
# Build graph
vae.model_fn(data=cifar)
# Training
vae.train(num_epochs=num_steps)
def main(args):
"""main procedure"""
# get configuration
device = global_conf["device"]
img_size = global_conf["image_size"]
data_dir = global_conf["data_dir"]
res_dir = global_conf["res_dir"]
# prepare data
train_loader, test_loader = prepare_data(args, dir_path=data_dir)
# prepare model
model = VAE(img_size[0] * img_size[1], args.n_hidden, args.dim_z,
img_size[0] * img_size[1])
optimizer = optim.Adam(model.parameters(), lr=args.lr)
# train and test
losses = []
for epoch in range(1, args.epochs + 1):
avg_loss = train(model, train_loader, epoch, optimizer, args, device,
img_size)
losses.append(avg_loss)
test(model, test_loader, epoch, args, device, img_size, res_dir)
with torch.no_grad():
sample = model.sample(64, device).cpu()
save_image(sample.view(64, 1, img_size[0], img_size[1]),
res_dir + '/sample_' + str(epoch) + '.png')
# plot train losses
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.plot(losses)
plt.savefig(res_dir + '/loss.png')
def predict():
with open("config.yml", 'r') as ymlfile:
params = yaml.safe_load(ymlfile)
model_path = '../logs/model-2019-08-22-162729/' + 'vae_weights.h5'
vae = VAE(params['image_shape'], params['latent_dim'])
vae.model.load_weights(model_path)
decoder = vae.model.get_layer('model_1')
# Use the decoder network to turn arbitrary latent space vectors into images and display a 2D manifold of the faces
n = 5 # figure with 15x15 faces
face_size_i = params['image_shape'][0]
face_size_j = params['image_shape'][1]
color_channel = params['image_shape'][2]
figure = np.zeros(shape=(face_size_i * n, face_size_j * n, color_channel))
# Linearly spaced coordinates on the unit square were transformed through the inverse CDF (ppf) of the Gaussian
# to produce values of the latent variables z, since the prior of the latent space is Gaussian
grid_x = norm.ppf(np.linspace(0.55, 0.85, n))
grid_y = norm.ppf(np.linspace(0.35, 0.95, n))
for i, yi in enumerate(grid_x):
for j, xi in enumerate(grid_y):
z_sample = np.array([[xi, yi]])
z_sample = np.tile(z_sample, params['batch_size']).reshape(
params['batch_size'], 2)
x_decoded = decoder.predict(z_sample,
batch_size=params['batch_size'])
face = x_decoded[0].reshape(face_size_i, face_size_j,
color_channel)
figure[i * face_size_i:(i + 1) * face_size_i,
j * face_size_j:(j + 1) * face_size_j] = face
plt.figure(figsize=(10, 10))
plt.imshow(figure, cmap='Greys_r')
plt.show()
def test_build_and_start_session(self):
model = VAE(
[28, 28], [1000, 500, 250], [250, 500, 28 * 28], 30, VAE.LossType.SIGMOID_CROSS_ENTROPY, 0.0005, 0.001
)
model.start_session()
model.stop_session()
def load_vaes(H):
vae = None
#vae = VAE(H)
#if H.restore_path:
# #logprint(f'Restoring vae from {H.restore_path}')
# print('Restoring vae from :', H.restore_path)
# restore_params(vae, H.restore_path, map_cpu=True, local_rank=None, mpi_size=None)
ema_vae = VAE(H)
if H.restore_ema_path:
#logprint(f'Restoring ema vae from {H.restore_ema_path}')
restore_params(ema_vae,
H.restore_ema_path,
map_cpu=True,
local_rank=None,
mpi_size=None)
elif (vae):
ema_vae.load_state_dict(vae.state_dict())
ema_vae.requires_grad_(False)
#vae = vae.cuda(H.local_rank)
ema_vae = ema_vae.cuda(H.local_rank)
#vae = DistributedDataParallel(vae, device_ids=[H.local_rank], output_device=H.local_rank)
#if len(list(vae.named_parameters())) != len(list(vae.parameters())):
# raise ValueError('Some params are not named. Please name all params.')
#total_params = 0
#for name, p in vae.named_parameters():
# total_params += np.prod(p.shape)
#print("Totat Params : ", total_params)
#logprint(total_params=total_params, readable=f'{total_params:,}')
return vae, ema_vae
def main():
if not os.path.exists("generated"):
os.makedirs("generated")
print "--- load START---"
mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)
print "--- load FINISH ---"
with tf.Session() as sess:
vae = VAE(sess,
batch_size=FLAGS.batch_size,
epoch=FLAGS.epoch,
hidden_size=FLAGS.hidden_size)
if FLAGS.train:
vae.train(mnist)
else:
vae.load("model.ckpt")
test = mnist.test.images
# batch_num = len(test) / FLAGS.batch_size
i = 0
for i in range(1):
vae.test(
i, test[i * FLAGS.batch_size:i * FLAGS.batch_size +
FLAGS.batch_size])
def main(args):
dim = 64
latent_size = 512
channels = 3
best_model_path = "./vae_lambda001.pth"
# should give from outside
output_path = args.output_path
# load model
model = VAE(d=dim, zsize=latent_size, channels=channels)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
model.load_state_dict(torch.load(best_model_path, map_location=device))
# 4. sample from normal distribution
seed = 5
np.random.seed(seed)
samples = np.random.randn(32, latent_size)
torch_samples = torch.FloatTensor(samples)
torch_samples = torch_samples.to(device)
result = model.decode(torch_samples)
# concat images
rows = []
for row in range(8):
rows.append(
torch.cat([img for img in result[row * 4:(row + 1) * 4]], axis=1))
ret = torch.cat(rows, axis=2)
img = ret.cpu().detach().numpy().transpose(1, 2, 0)
scipy.misc.imsave(output_path, img)
def divcolor():
out_dir, listdir, featslistdir = get_dirpaths(args)
batchsize = args.batchsize
hiddensize = args.hiddensize
nmix = args.nmix
data = colordata(\
os.path.join(out_dir, 'images'), \
listdir=listdir,\
featslistdir=featslistdir,
split='test')
nbatches = np.int_(np.floor(data.img_num/batchsize))
data_loader = DataLoader(dataset=data, num_workers=args.nthreads,\
batch_size=batchsize, shuffle=True, drop_last=True)
model_vae = VAE()
model_vae.cuda()
model_vae.load_state_dict(torch.load('%s/models/model_vae.pth' % (out_dir)))
model_vae.train(False)
model_mdn = MDN()
model_mdn.cuda()
model_mdn.load_state_dict(torch.load('%s/models/model_mdn.pth' % (out_dir)))
model_mdn.train(False)
for batch_idx, (batch, batch_recon_const, batch_weights, \
batch_recon_const_outres, batch_feats) in \
tqdm(enumerate(data_loader), total=nbatches):
input_feats = Variable(batch_feats).cuda()
mdn_gmm_params = model_mdn(input_feats)
gmm_mu, gmm_pi = get_gmm_coeffs(mdn_gmm_params)
gmm_pi = gmm_pi.view(-1, 1)
gmm_mu = gmm_mu.view(-1, hiddensize)
for j in range(batchsize):
batch_j = np.tile(batch[j, ...].numpy(), (batchsize, 1, 1, 1))
batch_recon_const_j = np.tile(batch_recon_const[j, ...].numpy(), (batchsize, 1, 1, 1))
batch_recon_const_outres_j = np.tile(batch_recon_const_outres[j, ...].numpy(), \
(batchsize, 1, 1, 1))
input_color = Variable(torch.from_numpy(batch_j)).cuda()
input_greylevel = Variable(torch.from_numpy(batch_recon_const_j)).cuda()
curr_mu = gmm_mu[j*nmix:(j+1)*nmix, :]
orderid = np.argsort(\
gmm_pi[j*nmix:(j+1)*nmix, 0].cpu().data.numpy().reshape(-1))
z = curr_mu.repeat(np.int((batchsize*1.)/nmix), 1)
_, _, color_out = model_vae(input_color, input_greylevel, z, is_train=False)
data.saveoutput_gt(color_out.cpu().data.numpy()[orderid, ...], \
batch_j[orderid, ...], \
'divcolor_%05d_%05d' % (batch_idx, j), \
nmix, \
net_recon_const=batch_recon_const_outres_j[orderid, ...])
def train_mdn(logger=None):
out_dir, listdir, featslistdir = get_dirpaths(args)
batchsize = args.batchsize
hiddensize = args.hiddensize
nmix = args.nmix
nepochs = args.epochs_mdn
data = colordata(\
os.path.join(out_dir, 'images'), \
listdir=listdir,\
featslistdir=featslistdir,
split='train')
nbatches = np.int_(np.floor(data.img_num/batchsize))
data_loader = DataLoader(dataset=data, num_workers=args.nthreads,\
batch_size=batchsize, shuffle=True, drop_last=True)
model_vae = VAE()
model_vae.cuda()
model_vae.load_state_dict(torch.load('%s/models/model_vae.pth' % (out_dir)))
model_vae.train(False)
model_mdn = MDN()
model_mdn.cuda()
model_mdn.train(True)
optimizer = optim.Adam(model_mdn.parameters(), lr=1e-3)
itr_idx = 0
for epochs_mdn in range(nepochs):
train_loss = 0.
for batch_idx, (batch, batch_recon_const, batch_weights, _, batch_feats) in \
tqdm(enumerate(data_loader), total=nbatches):
input_color = Variable(batch).cuda()
input_greylevel = Variable(batch_recon_const).cuda()
input_feats = Variable(batch_feats).cuda()
z = Variable(torch.randn(batchsize, hiddensize))
optimizer.zero_grad()
mu, logvar, _ = model_vae(input_color, input_greylevel, z)
mdn_gmm_params = model_mdn(input_feats)
loss, loss_l2 = mdn_loss(mdn_gmm_params, mu, torch.sqrt(torch.exp(logvar)), batchsize)
loss.backward()
optimizer.step()
train_loss = train_loss + loss.data[0]
if(logger):
logger.update_plot(itr_idx, [loss.data[0], loss_l2.data[0]], plot_type='mdn')
itr_idx += 1
train_loss = (train_loss*1.)/(nbatches)
print('[DEBUG] Training MDN, epoch %d has loss %f' % (epochs_mdn, train_loss))
torch.save(model_mdn.state_dict(), '%s/models/model_mdn.pth' % (out_dir))
def train_model_art(model_dir, num_steps, batch_size=64, learning_rate=0.0005):
"""
Train a VAE on the art data
:param model_dir:
:param num_steps:
:param batch_size:
:param learning_rate:
:return:
"""
with tf.Session() as sess:
vae = VAE(sess=sess,
model_dir=model_dir,
batch_size=batch_size,
learning_rate=learning_rate,
height=128,
width=128,
cdim=1,
n_z=128)
# Load cifar dataset
data = load_art_data(
data_path="/home/kevin/deep_learning/cat-dataset/cats/CAT_00",
batch_size=batch_size)
# Training
vae.train(data=data, num_epochs=num_steps)
def main():
# Load the HyperParameters
params = json.load(open(CONFIG_PATH))[0]
utils = ActionUtils(params['env_name'])
action_size = utils.action_size()
params['action_size'] = action_size
mdn_hps = params['mdn_hps']
mdn_hps['max_seq_len'] = params['max_seq_len']
mdn_hps['in_width'] = params['latent_size'] + action_size
mdn_hps['out_width'] = params['latent_size']
mdn_hps['action_size'] = action_size
mdn_hps['rnn_size'] = params['hidden_size']
mdn_hps = MDNRNN.set_hps_to_inference(mdn_hps)
# Create the MDN and load the params
mdnrnn = MDNRNN(mdn_hps)
mdnrnn.load(MDNRNN_PATH)
# Create the VAE
vae = VAE()
vae.make_vae_shape(params['img_size'], params['img_size'],
params['latent_size'])
vae.load_model(VAE_PATH)
# Create the Gym Env
env = gym.make(params['env_name'])
dream_vis(env, mdnrnn, vae, params, mdn_hps, "dream_1")
dream_vis(env, mdnrnn, vae, params, mdn_hps, "dream_2")
dream_vis(env, mdnrnn, vae, params, mdn_hps, "dream_3")
frame_vis(env, mdnrnn, vae, params, mdn_hps, "cmp_1")
frame_vis(env, mdnrnn, vae, params, mdn_hps, "cmp_2")
def train_model_mnist(model_dir,
num_steps,
batch_size=64,
learning_rate=0.0001):
"""
Train a VAE model
:param model_dir: model directory
:param num_steps: number of training steps
:param learning_rate: learning rate to use
:return:
"""
with tf.Session() as sess:
vae = VAE(sess=sess,
model_dir=model_dir,
batch_size=batch_size,
learning_rate=learning_rate,
height=28,
width=28,
cdim=1)
# Load data
mnist = load_mnist_data(batch_size=batch_size)
# Build graph
vae.model_fn(data=mnist)
# Training
vae.train(data="placeholder", num_epochs=num_steps)
def next_gen_process(id, model_info):
latent_size = 32
vae_model = VAE(latent_size)
checkpoint_dir = './vae_ckpt/'
latest = tf.train.latest_checkpoint(checkpoint_dir)
vae_model.load_weights(latest)
s = np.random.randint(1e8)
np.random.seed(s + id) # different seed for each process
env = gym.make('CarRacing-v0', verbose=0)
env = wrap_env(env, W, H, gray_scale=False)
results = []
for info in model_info:
m = Model()
m.set_weights(info['weights'])
m.copy_model(info['attr'], from_pickle=True)
res = m.evaluate(env, vae_model)
results.append(res)
env.close()
return results
def demo_embedding():
# get input data
mnist_data = input_data.read_data_sets('../../dataset/mnist_data',
one_hot=True)
num_sample = mnist_data.train.num_examples
batch_size = 100
network_architecture = \
dict(n_hidden_encoder_1=500, # 1st layer encoder neurons
n_hidden_encoder_2=500, # 2nd layer encoder neurons
n_hidden_decoder_1=500, # 1st layer decoder neurons
n_hidden_decoder_2=500, # 2nd layer decoder neurons
n_input=784, # MNIST data input (img shape: 28*28)
n_z=2) # dimensionality of latent space
# define model
vae_model = VAE(network_architecture, batch_size=batch_size)
# train the model
train(model=vae_model,
inputs=mnist_data.train,
num_epoch=50,
num_sample=num_sample,
batch_size=batch_size)
# embedding data into 2D space
X_sample, Y_sample = mnist_data.test.next_batch(5000)
z_mu = vae_model.embedding(X_sample)
plt.figure(figsize=(8, 6))
plt.scatter(z_mu[:, 0], z_mu[:, 1], c=np.argmax(Y_sample, 1))
plt.colorbar()
plt.grid()
plt.show()
def get_model(dataset, hps):
if dataset == 'mnist' or dataset == 'fashion':
model = VAE(hps)
elif dataset == 'cifar': # convolutional VAE for CIFAR
model = CVAE(hps)
return model
def __init__(self, args):
self.log_path = args.log_path
self.device = torch.device("cuda:0" if args.cuda else "cpu")
self.img_size = args.img_size
self.sample_num = args.sample_num
self.transform = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize([64, 64], 1),
transforms.ToTensor(),
transforms.Normalize([.5], [.5])
])
self.pil_transform = transforms.ToPILImage(mode="RGB")
self.norm_scale = np.loadtxt(os.path.join(args.config_path,
"norm_scale.txt"),
dtype=np.float32,
delimiter=",")[None]
self.norm_min = np.loadtxt(os.path.join(args.config_path,
"norm_min.txt"),
dtype=np.float32,
delimiter=",")[None]
self.pb_list = torch.from_numpy(
np.loadtxt(os.path.join(args.config_path, "pb_list.txt"),
dtype=np.float32,
delimiter=","))
self.kmeans = KMeans(n_clusters=2)
self.kmeans.fit(self.pb_list)
print("=" * 5, "Init LSTMPB", "=" * 5)
self.rnn = LSTMPB(args, pb_unit=self.pb_list[5][None])
pt_file = load_model(args.model_load_path, "*/*LSTMPB*.pt")
self.rnn.load_state_dict(torch.load(pt_file))
print("=" * 5, "Init VAE", "=" * 5)
self.vae = VAE(img_size=args.img_size, z_dim=args.vae_z_dims)
pt_file = load_model(args.model_load_path, "*/VAE*.pt")
self.vae.load_state_dict(torch.load(pt_file))
self.vae.eval()
print("=" * 5, "Init CVAE", "=" * 5)
self.cvae = CVAE(img_size=args.img_size, z_dim=args.cvae_z_dims)
pt_file = load_model(args.model_load_path, "*/*CVAE*.pt")
self.cvae.load_state_dict(torch.load(pt_file))
self.cvae.eval()
self.norm_mode = {
"joint": [0, 1, 2, 3, 4],
"visual": [5, 6, 7, 8, 9, 10, 11]
}
self.norm_mode[
"all"] = self.norm_mode["joint"] + self.norm_mode["visual"]
self.global_step = 0
self.his_log = HistoryWindow(maxlen=args.window_size)
#visualize current goal
_, goal = self.vae.decoder(self.denorm(self.goal, "visual"))
goal = ((goal[0] * .5 + .5) * 255).to(torch.int8)
self.goal_img = self.pil_transform(goal)
def gaussian_mnist(latent_dim=2, pixel_std=.05, k=1):
# SINGLE 'SAMPLES' VARIABLE
k_samples = K.variable(k, name='k_samples', dtype='int32')
# RECOGNITION MODEL
inpt = Input(shape=(mnist_pixels, ))
q_hidden_1 = Dense(64, activation='relu')(inpt)
q_hidden_2 = Dense(64, activation='relu')(q_hidden_1)
# LATENT -- PRIOR
latent = DiagonalGaussianLatent(dim=latent_dim,
prior=IsoGaussianPrior(latent_dim),
k_samples=k_samples)
latent_sample = latent(q_hidden_2)
# GENERATIVE MODEL
gen_hidden_1 = Dense(64, activation='relu')(latent_sample)
gen_hidden_2 = Dense(64, activation='relu')(gen_hidden_1)
reconstruction = Dense(mnist_pixels, activation='sigmoid')(gen_hidden_2)
# LIKELIHOOD
# Note: in some models, pixel_std is not constant but is also an output of the model so that it
# can indicate its own uncertainty.
likelihood = DiagonalGaussianLikelihood(reconstruction, pixel_std)
# Combine the above parts into a single model
return VAE(inpt=inpt,
latent=latent,
reconstruction=reconstruction,
likelihood=likelihood,
k_samples=k_samples)
def encode(_input, weights, latent, reconstruct):
conv_vae = VAE()
conv_vae.make_vae(_input + ".npz", int(latent))
conv_vae.load_model(weights)
latent_vectors = []
# for i in range(len(conv_vae.x_test)):
# latent_vectors.append(conv_vae.encode_image(conv_vae.x_test[i:i+1]))
raw_data = np.load(_input + ".npz")
for f in sorted(raw_data.files):
images = raw_data[f]
latent_vectors.append(
[conv_vae.encode_image(np.array([image]))[0] for image in images])
np.savez_compressed(_input + "_latent.npz", *latent_vectors)
if reconstruct:
data = np.load(_input + "_latent.npz")
# files = data.files
# vectors = np.array(data[files[0]])
reconstructed_images = []
for f in sorted(data.files):
latents = data[f]
# a = a.reshape(-1, a.shape[-2], a.shape[-1])
reconstructed_images.append(
[conv_vae.decode_latent(np.array([l]))[0] for l in latents])
# print(np.shape(np.array(reconstructed_images)))
# plt.imshow(reconstructed_images[1])
# plt.show()
np.savez_compressed(_input + "_recon.npz", *reconstructed_images)
def mnist_hvae(cls='vae', latent_dim=2, layers=3, pixel_std=.05, k=1):
# RECOGNITION MODEL(s)
q_model = Sequential()
q_model.add(Dense(64, activation='relu', input_dim=mnist_pixels))
q_model.add(Dense(64, activation='relu'))
# LATENT
latent = DiagonalGaussianLatent(dim=latent_dim)
# PRIOR
prior = IsoGaussianPrior(dim=latent_dim)
# GENERATIVE MODEL
p_model = Sequential()
p_model.add(Dense(64, activation='relu', input_dim=latent_dim))
p_model.add(Dense(64, activation='relu'))
p_model.add(Dense(mnist_pixels, activation='sigmoid'))
# LIKELIHOOD
likelihood = GaussianLikelihood(pixel_std)
# Combine the above parts into a single model
kwargs = {
'q_model': q_model,
'latent': latent,
'prior': prior,
'p_model': p_model,
'likelihood': likelihood,
'k': k
}
if cls == 'vae':
return VAE(**kwargs)
elif cls == 'iwae':
return IWAE(**kwargs)
def main(args):
""" main() driver function """
# Parameters parsing
if filepath_is_not_valid(args.config):
logging.error("The path {} is not a file. Aborting..".format(args.config))
exit()
configuration, architecture, hyperparameters = parse_config_file(args.config, args.variation)
dataset_info = prepare_dataset(configuration)
if (dataset_info is None):
exit()
# Initialization
model = None
if (args.variation == "VAE"):
model = VAE(architecture, hyperparameters, dataset_info)
elif (args.variation == "B-VAE"):
model = betaVAE(architecture, hyperparameters, dataset_info)
# here you can change the gpus parameter into the amount of gpus you want the model to use
trainer = Trainer(max_epochs = hyperparameters["epochs"], gpus=None, fast_dev_run=False)
# Training and testing
trainer.fit(model)
result = trainer.test(model)
# Model needs to be transferred to the cpu as sample and reconstruct are custom methods
model = model.cpu()
model.sample(5)
model.reconstruct(5)
def load_model(self, controller_weights=None):
p = self.params
self.action_utils = ActionUtils(p['env_name'])
self.action_size = self.action_utils.action_size()
self.vae = VAE()
sys.stdout = open(os.devnull, 'w')
self.vae.make_vae_shape(p['img_size'], p['img_size'], p['latent_size'])
sys.stdout = sys.__stdout__
self.vae.load_model('../' + p['vae_hps']['weights_path'])
# TODO: Make MDN just take in all of params.
mdn_hps = p['mdn_hps']
mdn_hps['max_seq_len'] = p['max_seq_len']
mdn_hps['in_width'] = p['latent_size'] + self.action_size
mdn_hps['out_width'] = p['latent_size']
mdn_hps['action_size'] = self.action_size
mdn_hps['rnn_size'] = p['hidden_size']
mdn_hps['batch_size'] = 1
mdn_hps['max_seq_len'] = 1
mdn_hps['use_recurrent_dropout'] = 0
mdn_hps['training'] = 0
# self.mdn_rnn = MDNRNN(mdn_hps)
# hps_inf = MDNRNN.set_hps_to_inference(hps)
self.mdn_rnn = MDNRNN(mdn_hps)
self.mdn_rnn.load('../' + p['mdn_hps']['weights_path'])
self.controller = ControllerModel(
[p['latent_size'] + p['hidden_size'], self.action_size])
if controller_weights:
self.controller.load_weights(controller_weights)
def main():
config = get_config()
logging.basicConfig(
format='%(asctime)s | %(message)s',
handlers=[
logging.FileHandler(os.path.join(config.log_root,
config.log_name)),
logging.StreamHandler()],
level=logging.INFO)
transformer = transforms.ToTensor()
train_dataset = datasets.MNIST(config.data_root, train=True, download=True, transform=transformer)
test_dataset = datasets.MNIST(config.data_root, train=False, download=True, transform=transformer)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True,
num_workers=4, pin_memory=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=config.batch_size, shuffle=False,
num_workers=4, pin_memory=True)
model = VAE()
trainer = Trainer(model, train_loader, test_loader, Adam(model.parameters(), lr=0.0002, betas=(0.5, 0.999)),
loss_function, 1)
trainer.train(config.epochs, config.log_metrics_every)
def main(*,
layers,
hidden_fn=nn.ReLU(),
model_name='VAE',
model_path='./vae.pt',
wd=1e-4,
num_epochs=30):
batch_size = 128
train_loader = Data.DataLoader(
dataset=JesterRatingDataset('../jester/train_ratings.csv'),
batch_size=batch_size,
shuffle=True)
print('已加载训练集数据...')
test_rating = JesterRatingDataset('../jester/test_ratings.csv')
print('\n已加载测试集数据...')
vae = VAE(*layers, hidden_fn=hidden_fn)
train_vae_model(vae,
train_loader,
test_rating=test_rating.ratings,
model_name=model_name,
model_path=model_path,
wd=wd,
num_epochs=num_epochs)