def main():
parser = argparse.ArgumentParser()
parser.add_argument('--train', help='train variational autoencoder', action='store_true')
parser.add_argument('--generate',help='generate images', action='store_true')
parser.add_argument('--traverse',help='traverse latent space', action='store_true')
args = parser.parse_args()
sess = make_session()
if(args.train):
#prepare_tfrecords()
training_dataset = prepare_dataset(Config.tfrecord_dir+'train.tfrecord')
iterator = tf.compat.v1.data.make_one_shot_iterator(training_dataset)
vae = VAE(sess, training=True, training_iterator = iterator)
train(sess, vae)
if(args.generate) or (args.traverse):
vae = VAE(sess, training=False, restore=True)
if(args.generate):
random_vector = tf.random.normal(shape=[Config.num_gen_imgs, vae.latent_dim])
filename = Config.results_dir + 'randomFakeGrid.jpg'
gen_imgs = generate_fake_images(sess, vae, random_vector, filename)
if(args.traverse):
traverse_latents(sess, vae)
def main(args):
# Check if the output folder is exist
if not os.path.exists('./vae_results/'):
os.mkdir('./vae_results/')
# Load data
torch.manual_seed(args.seed)
kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}
train_loader = torch.utils.data.DataLoader(
datasets.MNIST('./data', train=True, download=True,
transform=transforms.ToTensor()),
batch_size=args.batch_size, shuffle=True, **kwargs)
# Load model
model = VAE().cuda() if torch.cuda.is_available() else VAE()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# Train and generate sample every epoch
loss_list = []
for epoch in range(1, args.epochs + 1):
model.train()
_loss = train(epoch, model, train_loader, optimizer)
loss_list.append(_loss)
model.eval()
sample = torch.randn(64, 20)
sample = Variable(sample).cuda() if torch.cuda.is_available() else Variable(sample)
sample = model.decode(sample).cpu()
save_image(sample.view(64, 1, 28, 28).data, 'vae_results/sample_' + str(epoch) + '.png')
plt.plot(range(len(loss_list)), loss_list, '-o')
plt.savefig('vae_results/vae_loss_curve.png')
def train(data_loader, model_index, x_eval_train, loaded_model):
### Model Initiation
if loaded_model:
vae = VAE()
vae.cuda()
saved_state_dict = tor.load(loaded_model)
vae.load_state_dict(saved_state_dict)
vae.cuda()
else:
vae = VAE()
vae = vae.cuda()
loss_func = tor.nn.MSELoss().cuda()
#optim = tor.optim.SGD(fcn.parameters(), lr=LR, momentum=MOMENTUM)
optim = tor.optim.Adam(vae.parameters(), lr=LR)
lr_step = StepLR(optim, step_size=LR_STEPSIZE, gamma=LR_GAMMA)
### Training
for epoch in range(EPOCH):
print("|Epoch: {:>4} |".format(epoch + 1))
### Training
for step, (x_batch, y_batch) in enumerate(data_loader):
print("Process: {}/{}".format(step,
int(AVAILABLE_SIZE[0] / BATCHSIZE)),
end="\r")
x = Variable(x_batch).cuda()
y = Variable(y_batch).cuda()
out, KLD = vae(x)
recon_loss = loss_func(out.cuda(), y)
loss = (recon_loss + KLD_LAMBDA * KLD)
loss.backward()
optim.step()
lr_step.step()
optim.zero_grad()
if step % RECORD_JSON_PERIOD == 0:
save_record(model_index, epoch, optim, recon_loss, KLD)
if step % RECORD_PIC_PERIOD == 0:
save_pic("output_{}".format(model_index), vaee, 3)
if step % RECORD_MODEL_PERIOD == 0:
tor.save(
vae.state_dict(),
os.path.join(MODEL_ROOT,
"ave_model_{}.pkl".format(model_index)))
def test_vae(self):
dae = DAE()
vae = VAE(dae)
vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, "vae")
# Check size of optimizing vars
self.assertEqual(len(vars), 14+12)
def train_and_test():
# pick a test sample
dataset = Data(config.train_size, config.batch_size, config.test_size)
for test_batch in dataset.test_dataset:
test_sample = test_batch[:config.num_examples_to_generate, :, :, :]
random_vector_for_generation = tf.random.normal(
shape=[config.num_examples_to_generate, config.latent_dim])
model = VAE(config.latent_dim)
#utils.generate_and_save_images(model, 0, test_sample)
optimizer = tf.keras.optimizers.Adam(1e-4)
for epoch in range(1, config.epochs + 1):
start_time = time.time()
for train_x in dataset.train_dataset:
utils.train_step(model, train_x, optimizer=optimizer)
end_time = time.time()
loss = tf.keras.metrics.Mean()
for test_x in dataset.test_dataset:
loss(utils.compute_loss(model, test_x))
elbo = -loss.result()
print(
'Epoch: {}, Test set ELBO: {}, time elapse for current epoch: {}'.
format(epoch, elbo, end_time - start_time))
utils.generate_and_save_images(model, epoch, test_sample)
def train_VAE():
model = VAE()
model.summary()
history = []
x_train = import_image(folder='train')
x_test = import_image(folder='test')
#x_train = import_image('test')
arbitrary = np.zeros([x_train.shape[0], 1024 * 2])
model.summary()
history = model.fit(x_train, [x_train, arbitrary],
epochs=200,
batch_size=128,
shuffle=True,
verbose=1,
validation_data=(x_test,
[x_test,
arbitrary[:x_test.shape[0]]]))
print("history =", history.history.keys())
output_loss1 = history.history['recons_loss']
output_loss2 = history.history['KLD_loss']
np.save('VAE2_recons_loss', output_loss1)
np.save('VAE2_KLD_loss', output_loss2)
model.save_weights('saved_model/VAE2_epochs60_weights.h5')
def train():
model = VAE()
optimizer = tf.train.AdamOptimizer(learning_rate)
train_step = optimizer.minimize(model.loss)
saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
x = mnist.train.images
n_iters = int(np.ceil(len(x) / batch_size))
for epoch in range(n_epochs):
for i in trange(n_iters, desc=f'Epoch {epoch}'):
_, loss = sess.run([train_step, model.loss],
feed_dict={
model.x: x[i::n_iters],
})
if i % 100 == 0:
tqdm.write(f'Step {i}, loss: {loss}')
print(saver.save(sess, model_path))
test_image = mnist.test.images[:1]
out = sess.run(model.out, feed_dict={model.x: test_image})
out_image = np.squeeze(out[0]) * 255
cv2.imwrite(f'samples/{epoch}_{int(loss)}.png', out_image)
def train():
use_cuda = True
epochs = 100
batch_size = 32
hidden_dim = 500
z_dim = 20
lr = 0.0001
compressed = transforms.Compose([dataset.ToTensor()])
variable = dataset.ToVariable(use_cuda=use_cuda)
kwargs = {'num_workers': 8, 'pin_memory': True}
loaders = dataset.setup_data_loaders(dataset.LastFMCSVDataset,
use_cuda,
batch_size,
transform=compressed,
**kwargs)
print('{} steps for all data / 1 epoch'.format(len(loaders['X'])))
vae = VAE(1, hidden_dim, z_dim, use_cuda=use_cuda)
adagrad_param = {'lr': lr}
optimizer = optim.Adam(vae.parameters(), **adagrad_param)
for epoch in range(epochs):
loss = train_epoch(vae, optimizer, loaders, variable)
print('Epoch{}:{}'.format(epoch, loss))
def train(args):
data_path = args.data
data = pd.read_csv(data_path, index_col=0)
data = prepare_data(data)
batch_size = args.batch_size
lr = args.lr
n_epochs = args.n_epochs
verbose = args.verbose
model = VAE(input_size=data.shape[0], batch_size=batch_size)
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
criterion = nn.MSELoss()
model.train()
for epoch in range(n_epochs):
loss_sum = 0
for i in range(data.shape[1] // batch_size):
batch = Variable(data[:, i * batch_size:(i + 1) * batch_size],
requires_grad=False)
optimizer.zero_grad()
recon_batch = model(batch)
loss = criterion(recon_batch, batch)
loss.backward()
optimizer.step()
loss_sum += loss.data.numpy()
if verbose:
print('Epoch: {} - loss: {:.2f}'.format(epoch, loss_sum))
def slave(comm):
mus, logvars = load_init_z()
vae = VAE()
vae.load_state_dict(torch.load(cfg.vae_save_ckpt, map_location=lambda storage, loc: storage)['model'])
model = RNNModel()
model.load_state_dict(torch.load(cfg.rnn_save_ckpt, map_location=lambda storage, loc: storage)['model'])
count = 1
status = MPI.Status()
gpuid = comm.rank % 4
# device = torch.device('cuda:{}'.format(gpuid))
# vae.to(device)
# model.to(device)
print('Worker {} Started, model on GPU {}'.format(comm.rank, gpuid))
while True:
solution = comm.recv(source=0, tag=MPI.ANY_TAG, status=status)
tag = status.Get_tag()
if tag == 1:
print('Worker {} received solution {}'.format(comm.rank, count))
zs = [sample_init_z(mus, logvars) for _ in range(cfg.trials_per_pop)]
controller = deflatten_controller(solution)
reward = rollout(model, controller, zs)
print('Worker {} finished solution {}, reward: mean {} | max {} | min {} | std {}'.format(
comm.rank, count, reward.mean(), reward.max(), reward.min(), reward.std()))
comm.send(reward.mean(), dest=0, tag=2)
count += 1
elif tag == 3:
print('Worker {} evaluate current solution'.format(comm.rank))
controller = deflatten_controller(solution)
reward = evaluate(model, vae, controller)
comm.send(reward, dest=0, tag=2)
def run(args):
train_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
args.data, transform=transforms.ToTensor()),
batch_size=args.batch_size,
shuffle=True,
num_workers=8)
test_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
args.data_val, transform=transforms.ToTensor()),
batch_size=args.batch_size,
shuffle=False,
num_workers=8)
model = nn.DataParallel(VAE())
model = model.to(args.device)
if args.checkpoint is not None:
model.load_state_dict(torch.load(args.checkpoint))
optimizer = optim.Adam(model.parameters(), lr=args.lr)
for epoch in range(1, args.epochs + 1):
train(epoch, model, train_loader, optimizer, args)
model_file = 'models/model_' + str(epoch) + '.pth'
torch.save(model.state_dict(), model_file)
print('Saved model to ' + model_file)
test(epoch, model, test_loader, args)
writer.close()
def __init__(self,
hparams,
train_size: int,
class_weight: Optional[Tensor] = None):
# model, criterion
self.model = VAE()
# optimizer and scheduler
self.optimizer = torch.optim.Adam(self.model.parameters(),
lr=hparams.learning_rate,
eps=hparams.eps,
weight_decay=hparams.weight_decay)
self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
self.optimizer, **hparams.scheduler)
self.bce = nn.BCEWithLogitsLoss(reduction='none')
# self.kld = nn.KLDivLoss(reduction='sum')
# device
device_for_summary = self.__init_device(hparams.device,
hparams.out_device)
# summary
self.writer = SummaryWriter(logdir=hparams.logdir)
# TODO: fill in ~~DUMMY~~INPUT~~SIZE~~
path_summary = Path(self.writer.logdir, 'summary.txt')
if not path_summary.exists():
print_to_file(path_summary, summary, (self.model, (40, 11)),
dict(device=device_for_summary))
# save hyperparameters
path_hparam = Path(self.writer.logdir, 'hparams.txt')
if not path_hparam.exists():
print_to_file(path_hparam, hparams.print_params)
def run(config, seed, device):
set_global_seeds(seed)
logdir = Path(config['log.dir']) / str(config['ID']) / str(seed)
train_loader, test_loader = make_dataset(config)
if config['nn.type'] == 'VAE':
model = VAE(config, device)
elif config['nn.type'] == 'ConvVAE':
model = ConvVAE(config, device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
engine = Engine(config,
model=model,
optimizer=optimizer,
train_loader=train_loader,
test_loader=test_loader)
train_logs = []
eval_logs = []
for epoch in range(config['train.num_epoch']):
train_logger = engine.train(epoch, logdir=logdir)
train_logs.append(train_logger.logs)
eval_logger = engine.eval(epoch, logdir=logdir)
eval_logs.append(eval_logger.logs)
pickle_dump(obj=train_logs, f=logdir / 'train_logs', ext='.pkl')
pickle_dump(obj=eval_logs, f=logdir / 'eval_logs', ext='.pkl')
return None
def extract(fs, idx, N):
model = VAE()
model.load_state_dict(
torch.load(cfg.vae_save_ckpt,
map_location=lambda storage, loc: storage)['model'])
model = model.cuda(idx)
for n, f in enumerate(fs):
data = np.load(f)
imgs = data['sx'].transpose(0, 3, 1, 2)
actions = data['ax']
rewards = data['rx']
dones = data['dx']
x = torch.from_numpy(imgs).float().cuda(idx) / 255.0
mu, logvar, _, z = model(x)
save_path = "{}/{}".format(cfg.seq_extract_dir, f.split('/')[-1])
np.savez_compressed(save_path,
mu=mu.detach().cpu().numpy(),
logvar=logvar.detach().cpu().numpy(),
dones=dones,
rewards=rewards,
actions=actions)
if n % 10 == 0:
print('Process %d: %5d / %5d' % (idx, n, N))
def train_vae():
model = VAE(h,w,256,50,device)
optimizer = optim.Adadelta(model.parameters())
l2 = lambda epoch: pow((1.-1.*epoch/epochs),0.9)
scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=l2)
for epoch in range(epochs):
scheduler.step()
model.train()
tr_recon_loss = 0
for batch_idx, (data, target) in enumerate(tr):
if batch_idx >= tr_size:
break
data = data.to(device)
optimizer.zero_grad()
recon_batch, zmu, zvar,_= model(data)
recon_loss, kl = loss_V(recon_batch, data, zmu,torch.exp(0.5*zvar))
loss = recon_loss + kl
loss.backward()
tr_recon_loss += recon_loss.item()
optimizer.step()
if batch_idx % log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tReconstruction-Loss: {:.4f}, KL: {:.4f}'.format(
epoch, batch_idx * len(data), len(mnist_tr),
100. * batch_idx / len(tr),
recon_loss / len(data),kl/len(data)))
print('====> Epoch: {} Reconstruction loss: {:.4f}'.format(
epoch, tr_recon_loss / (tr_size*mb_size)))
test(epoch,model)
return model
def test_real(epi):
vae = VAE()
vae.load_state_dict(torch.load(cfg.vae_save_ckpt)['model'])
model = RNNModel()
model.load_state_dict(torch.load(cfg.rnn_save_ckpt)['model'])
controller = Controller()
controller.load_state_dict(torch.load(cfg.ctrl_save_ckpt)['model'])
env = DoomTakeCover(True)
obs = env.reset()
model.reset()
frames = []
for step in range(cfg.max_steps):
frames.append(cv2.resize(obs, (256, 256)))
obs = torch.from_numpy(obs.transpose(2, 0,
1)).unsqueeze(0).float() / 255.0
mu, logvar, _, z = vae(obs)
inp = torch.cat((model.hx.detach(), model.cx.detach(), z), dim=1)
y = controller(inp)
y = y.item()
action = encode_action(y)
model.step(z.unsqueeze(0), action.unsqueeze(0))
obs_next, reward, done, _ = env.step(action.item())
obs = obs_next
if done:
break
print('Episode {}: Real Reward {}'.format(epi, step))
write_video(frames, 'real_{}.avi'.format(epi), (256, 256))
os.system('mv real_{}.avi /home/bzhou/Dropbox/share'.format(epi))
def train():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
vae = VAE(device)
vae.to(device)
optimizer = optim.SGD(vae.parameters(), lr=lr)
#viz = visdom.Visdom()
print(time.asctime(time.localtime(time.time())))
n = 0
for epoch in range(epoch_sum):
loss_sum = 0
for idx, image in enumerate(train_dataloader, 0):
image = image[0].to(device)
optimizer.zero_grad()
out, mu, logvar = vae(image)
loss = loss_function(out, image, mu, logvar)
loss_sum += loss.item()
loss.backward()
optimizer.step()
#viz.line([loss.item()], [n], update="append", win='loss_win')
n += 1
if idx % 100 == 99:
print("epoch:%d, idx:%d, loss:%.6f" %
(epoch + 1, idx, loss_sum / 100))
loss_sum = 0
torch.save(vae.state_dict(), './vae.pth')
print(time.asctime(time.localtime(time.time())))
def load_model_and_dataset(checkpt_filename):
checkpt = torch.load(checkpt_filename)
args = checkpt['args']
state_dict = checkpt['state_dict']
# backwards compatibility
if not hasattr(args, 'conv'):
args.conv = False
from model import VAE, setup_data_loaders
# model
prior_dist = dist.Normal()
q_dist = dist.Normal()
vae = VAE(z_dim=args.latent_dim,
use_cuda=False,
prior_dist=prior_dist,
q_dist=q_dist,
conv=args.conv)
vae.load_state_dict(state_dict, strict=False)
vae.eval()
# dataset loader
loader = setup_data_loaders(args, use_cuda=False)
return vae, loader
def main(argv):
manager = DataManager()
manager.load()
sess = tf.Session()
#sess = tf_debug.LocalCLIDebugWrapperSession(sess)
model = VAE(beta=flags.beta,
learning_rate=flags.learning_rate,
alpha = flags.alpha,
kappa = flags.kappa,
lagrange_mult_param = flags.lagrange_mult_param,
use_geco = flags.use_geco)
sess.run(tf.global_variables_initializer())
saver = load_checkpoints(sess)
if flags.training:
# Train
train(sess, model, manager, saver)
else:
reconstruct_check_images = manager.get_random_images(10)
# Image reconstruction check
reconstruct_check(sess, model, reconstruct_check_images)
# Disentangle check
disentangle_check(sess, model, manager)
def __init__(self, args):
self.args = args
torch.manual_seed(self.args.seed)
if self.args.cuda:
torch.cuda.manual_seed(self.args.seed)
kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}
train_loader = torch.utils.data.DataLoader(
datasets.MNIST('./data',
train=True,
download=True,
transform=transforms.ToTensor()),
batch_size=self.args.batch_size,
shuffle=True,
**kwargs)
test_loader = torch.utils.data.DataLoader(
datasets.MNIST('./data',
train=False,
transform=transforms.ToTensor()),
batch_size=self.args.batch_size,
shuffle=True,
**kwargs)
self.train_loader = train_loader
self.test_loader = test_loader
self.model = VAE()
if self.args.cuda:
self.model.cuda()
self.optimizer = optim.Adam(self.model.parameters(), lr=1e-3)
def test_vae(self):
dae = DAE()
vae = VAE(dae)
vars = vae.get_vars()
# Check size of optimizing vars
self.assertEqual(len(vars), 14 + 12)
def main(args):
#
save_dir = os.path.join(args.save_dir, args.model_type)
img_dir = os.path.join(args.img_dir, args.model_type)
log_dir = os.path.join(args.log_dir, args.model_type)
train_dir = args.train_dir
if not os.path.exists(save_dir):
os.makedirs(save_dir)
if not os.path.exists(log_dir):
os.makedirs(log_dir)
if not os.path.exists(img_dir):
os.makedirs(img_dir)
mnist = utils.read_data_sets(args.train_dir)
summary_writer = tf.summary.FileWriter(log_dir)
config_proto = utils.get_config_proto()
sess = tf.Session(config=config_proto)
if args.model_type == "vae":
model = VAE(args, sess, name="vae")
elif args.model_type == "dcvae":
model = DCVAE(args, sess, name="dcvae")
total_batch = mnist.train.num_examples // args.batch_size
for epoch in range(1, args.nb_epoch + 1):
print "Epoch %d start with learning rate %f" % (
epoch, model.learning_rate.eval(sess))
print "- " * 50
epoch_start_time = time.time()
step_start_time = epoch_start_time
for i in range(1, total_batch + 1):
x_batch, y_batch = mnist.train.next_batch(args.batch_size)
_, loss, loss_rec, loss_kl, global_step, summaries = model.train(
x_batch)
summary_writer.add_summary(summaries, global_step)
step_start_time = time.time()
if global_step % args.log_period == 0:
print "global step %d, loss %.9f, loss_rec %.9f, loss_kl %.9f, time %.2fs" \
% (global_step, loss, loss_rec, loss_kl, time.time() - step_start_time)
step_start_time = time.time()
if args.anneal and epoch >= args.anneal_start:
sess.run(model.learning_rate_decay_op)
if epoch % args.save_period == 0:
z = np.random.normal(size=[100, args.latent_dim])
if args.model_type == "vae":
gen_images = np.reshape(model.generate(z), (100, 28, 28, 1))
elif args.model_type == "dcvae":
gen_images = np.reshape(model.generate(z, 100),
(100, 28, 28, 1))
utils.save_images(gen_images, [10, 10],
os.path.join(img_dir, "sample%s.jpg" % epoch))
model.saver.save(sess, os.path.join(save_dir, "model.ckpt"))
print "Model stored...."
def test_scan(self):
dae = DAE()
vae = VAE(dae)
scan = SCAN(dae, vae)
vars = scan.get_vars()
# Check size of optimizing vars
self.assertEqual(len(vars), 6 + 4)
def test_scan(self):
dae = DAE()
vae = VAE(dae)
scan = SCAN(dae, vae)
vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, "scan")
# Check size of optimizing vars
self.assertEqual(len(vars), 6+4)
def load_model(path):
restore_dict = torch.load(path)
model = VAE(**restore_dict["model"])
model.load_state_dict(restore_dict["model_state_dict"])
model.eval()
return model
def __init__(self, dataset):
self.model = VAE()
if config.USE_GPU:
self.model.cuda()
self, dataset = dataset
self.optimizer = optim.Adam(self.model.parameters(), lr=1e-3)
self.train_loader = DataLoader(self.dataset(train=True))
self.test_loader = DataLoader(self.dataset(train=False))
def main() -> None:
tokenizer = Tokenizer(args.vocab_file)
vocabulary_size = len(tokenizer)
dataset = SentenceDataset(args.input_file, tokenizer=tokenizer.encode)
loader = DataLoader(dataset,
args.batch_size,
shuffle=False,
collate_fn=dataset.collate_fn,
drop_last=False)
searcher = BeamSearch(tokenizer.eos_index, beam_size=args.search_width)
model = VAE(
num_embeddings=len(tokenizer),
dim_embedding=args.dim_embedding,
dim_hidden=args.dim_hidden,
dim_latent=args.dim_latent,
num_layers=args.num_layers,
bidirectional=args.bidirectional,
dropout=0.,
word_dropout=0.,
dropped_index=tokenizer.unk_index,
).to(device)
model.load_state_dict(torch.load(args.checkpoint_file,
map_location=device))
model.eval()
print('Generating sentence...')
all_hypotheses = []
with torch.no_grad():
for s in tqdm(loader):
s = s.to(device)
length = torch.sum(s != tokenizer.pad_index, dim=-1)
bsz = s.shape[0]
mean, logvar = model.encode(s, length)
# z = model.reparameterize(mean, logvar)
z = mean
hidden = model.fc_hidden(z)
hidden = hidden.view(bsz, -1,
model.dim_hidden).transpose(0,
1).contiguous()
start_predictions = torch.zeros(bsz, device=device).fill_(
tokenizer.bos_index).long()
start_state = {'hidden': hidden.permute(1, 0, 2)}
predictions, log_probabilities = searcher.search(
start_predictions, start_state, model.step)
for preds in predictions:
tokens = preds[0]
tokens = tokens[tokens != tokenizer.eos_index].tolist()
all_hypotheses.append(tokenizer.decode(tokens))
print('Done')
with open(args.output_file, 'w') as f:
f.write('\n'.join(all_hypotheses))
def main(args):
### VAE on MNIST
n_transform = transforms.Compose([transforms.ToTensor()])
dataset = MNIST('data', transform=n_transform, download=True)
### CVAE on MNIST
# n_transform = transforms.Compose([transforms.ToTensor()])
# dataset = MNIST('data', transform=n_transform)
### CVAE on facescrub-5
# n_transform = transforms.Compose([transforms.Resize(32), transforms.ToTensor()])
# dataset = ImageFolder('facescrub-5', transform=n_transform)
data_loader = DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=2)
vae = VAE(args.latent_size).cuda()
### CVAE
# vae = CVAE(args.latent_size, num_labels=args.num_labels).cuda()
optimizer = torch.optim.Adam(vae.parameters(), lr=args.learning_rate)
# Decide if you want to use fixed noise or not
fix_noise = Variable(torch.randn((50, args.latent_size)).cuda())
num_iter = 0
for epoch in range(args.epochs * 10):
for _, batch in enumerate(data_loader, 0):
img = Variable(batch[0].cuda())
label = Variable(batch[1].cuda())
recon_img, mean, log_var, z = vae(img)
### CVAE
# recon_img, mean, log_var, z = vae(img, label)
loss = loss_fn(recon_img, img, mean, log_var)
optimizer.zero_grad()
loss.backward()
optimizer.step()
num_iter += 1
if num_iter % args.print_every == 0:
print("Batch %04d/%i, Loss %9.4f" %
(num_iter, len(data_loader) - 1, loss.data.item()))
if num_iter % args.save_test_sample == 0:
x = vae.inference(fix_noise)
save_img(args, x.detach(), num_iter)
if num_iter % args.save_recon_img == 0:
save_img(args, recon_img.detach(), num_iter, recon=True)
def test_prepare_loss(self):
model = VAE()
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
self.assertEqual(model.loss.get_shape(), ())
self.assertEqual(model.x_out_logit.get_shape()[1], 64 * 64)
self.assertEqual(model.x_out.get_shape()[1], 64 * 64)
def main(args):
print("Loading data")
dataset = args.data.rstrip('/').split('/')[-1]
torch.cuda.set_device(args.cuda)
device = args.device
if dataset == 'mnist':
train_loader, test_loader = get_mnist(args.batch_size, 'data/mnist')
num = 10
elif dataset == 'fashion':
train_loader, test_loader = get_fashion_mnist(args.batch_size,
'data/fashion')
num = 10
elif dataset == 'svhn':
train_loader, test_loader, _ = get_svhn(args.batch_size, 'data/svhn')
num = 10
elif dataset == 'stl':
train_loader, test_loader, _ = get_stl10(args.batch_size, 'data/stl10')
elif dataset == 'cifar':
train_loader, test_loader = get_cifar(args.batch_size, 'data/cifar')
num = 10
elif dataset == 'chair':
train_loader, test_loader = get_chair(args.batch_size,
'~/data/rendered_chairs')
num = 1393
elif dataset == 'yale':
train_loader, test_loader = get_yale(args.batch_size, 'data/yale')
num = 38
model = VAE(28 * 28, args.code_dim, args.batch_size, 10,
dataset).to(device)
phi = nn.Sequential(
nn.Linear(args.code_dim, args.phi_dim),
nn.LeakyReLU(0.2, True),
).to(device)
model.load_state_dict(torch.load(args.fname))
if args.tsne:
datas, targets = [], []
for i, (data, target) in enumerate(test_loader):
datas.append(data), targets.append(target)
if i >= 5:
break
data, target = torch.cat(datas, dim=0), torch.cat(targets, dim=0)
c = F.one_hot(target.long(), num_classes=num).float()
_, _, _, z = model(data.to(args.device), c.to(args.device))
z, target = z.detach().cpu().numpy(), target.cpu().numpy()
tsne = TSNE(n_components=2, init='pca', random_state=0)
z_2d = tsne.fit_transform(z)
plt.figure(figsize=(6, 5))
for a in range(8):
for b in range(a + 1, 10):
plot_embedding(
z_2d,
target,
a,
b,
)
plt.savefig('tsne_c{}_{}_{}{}.png'.format(
int(args.c), dataset, a, b))