def main(args):
if args.save_path == '':
make_savepath(args)
seed(args)
if args.cuda:
print('using cuda')
print(args)
device = torch.device("cuda" if args.cuda else "cpu")
args.device = device
opt_dict = {"not_improved": 0, "lr": 1., "best_loss": 1e4}
all_data = torch.load(args.data_file)
x_train, x_val, x_test = all_data
x_train = x_train.to(device)
x_val = x_val.to(device)
x_test = x_test.to(device)
y_size = 1
y_train = x_train.new_zeros(x_train.size(0), y_size)
y_val = x_train.new_zeros(x_val.size(0), y_size)
y_test = x_train.new_zeros(x_test.size(0), y_size)
print(torch.__version__)
train_data = torch.utils.data.TensorDataset(x_train, y_train)
val_data = torch.utils.data.TensorDataset(x_val, y_val)
test_data = torch.utils.data.TensorDataset(x_test, y_test)
train_loader = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size,
shuffle=True)
val_loader = torch.utils.data.DataLoader(val_data,
batch_size=args.batch_size,
shuffle=True)
test_loader = torch.utils.data.DataLoader(test_data,
batch_size=args.batch_size,
shuffle=True)
print('Train data: %d batches' % len(train_loader))
print('Val data: %d batches' % len(val_loader))
print('Test data: %d batches' % len(test_loader))
sys.stdout.flush()
log_niter = len(train_loader) // 5
encoder = ResNetEncoderV2(args)
decoder = PixelCNNDecoderV2(args)
vae = VAE(encoder, decoder, args).to(device)
if args.sample_from != '':
save_dir = "samples/%s" % args.dataset
if not os.path.exists(save_dir):
os.makedirs(save_dir)
vae.load_state_dict(torch.load(args.sample_from))
vae.eval()
with torch.no_grad():
sample_z = vae.sample_from_prior(400).to(device)
sample_x, sample_probs = vae.decode(sample_z, False)
image_file = 'sample_binary_from_%s.png' % (
args.sample_from.split('/')[-1][:-3])
save_image(sample_x.data.cpu(),
os.path.join(save_dir, image_file),
nrow=20)
image_file = 'sample_cont_from_%s.png' % (
args.sample_from.split('/')[-1][:-3])
save_image(sample_probs.data.cpu(),
os.path.join(save_dir, image_file),
nrow=20)
return
if args.eval:
print('begin evaluation')
test_loader = torch.utils.data.DataLoader(test_data,
batch_size=50,
shuffle=True)
vae.load_state_dict(torch.load(args.load_path))
vae.eval()
with torch.no_grad():
test(vae, test_loader, "TEST", args)
au, au_var = calc_au(vae, test_loader)
print("%d active units" % au)
# print(au_var)
calc_iwnll(vae, test_loader, args)
return
enc_optimizer = optim.Adam(vae.encoder.parameters(), lr=0.001)
dec_optimizer = optim.Adam(vae.decoder.parameters(), lr=0.001)
opt_dict['lr'] = 0.001
iter_ = 0
best_loss = 1e4
best_kl = best_nll = best_ppl = 0
decay_cnt = pre_mi = best_mi = mi_not_improved = 0
aggressive_flag = True if args.aggressive else False
vae.train()
start = time.time()
kl_weight = args.kl_start
anneal_rate = (1.0 - args.kl_start) / (args.warm_up * len(train_loader))
for epoch in range(args.epochs):
report_kl_loss = report_rec_loss = 0
report_num_examples = 0
for datum in train_loader:
batch_data, _ = datum
batch_data = torch.bernoulli(batch_data)
batch_size = batch_data.size(0)
report_num_examples += batch_size
# kl_weight = 1.0
kl_weight = min(1.0, kl_weight + anneal_rate)
sub_iter = 1
batch_data_enc = batch_data
burn_num_examples = 0
burn_pre_loss = 1e4
burn_cur_loss = 0
while aggressive_flag and sub_iter < 100:
enc_optimizer.zero_grad()
dec_optimizer.zero_grad()
burn_num_examples += batch_data_enc.size(0)
loss, loss_rc, loss_kl = vae.loss(batch_data_enc,
kl_weight,
nsamples=args.nsamples)
burn_cur_loss += loss.sum().item()
loss = loss.mean(dim=-1)
loss.backward()
torch.nn.utils.clip_grad_norm_(vae.parameters(), clip_grad)
enc_optimizer.step()
id_ = np.random.choice(x_train.size(0),
args.batch_size,
replace=False)
batch_data_enc = torch.bernoulli(x_train[id_])
if sub_iter % 10 == 0:
burn_cur_loss = burn_cur_loss / burn_num_examples
if burn_pre_loss - burn_cur_loss < 0:
break
burn_pre_loss = burn_cur_loss
burn_cur_loss = burn_num_examples = 0
sub_iter += 1
# print(sub_iter)
enc_optimizer.zero_grad()
dec_optimizer.zero_grad()
loss, loss_rc, loss_kl = vae.loss(batch_data,
kl_weight,
nsamples=args.nsamples)
loss = loss.mean(dim=-1)
loss.backward()
torch.nn.utils.clip_grad_norm_(vae.parameters(), clip_grad)
loss_rc = loss_rc.sum()
loss_kl = loss_kl.sum()
if not aggressive_flag:
enc_optimizer.step()
dec_optimizer.step()
report_rec_loss += loss_rc.item()
report_kl_loss += loss_kl.item()
if iter_ % log_niter == 0:
train_loss = (report_rec_loss +
report_kl_loss) / report_num_examples
if aggressive_flag or epoch == 0:
vae.eval()
with torch.no_grad():
mi = calc_mi(vae, val_loader)
au, _ = calc_au(vae, val_loader)
vae.train()
print('epoch: %d, iter: %d, avg_loss: %.4f, kl: %.4f, mi: %.4f, recon: %.4f,' \
'au %d, time elapsed %.2fs' %
(epoch, iter_, train_loss, report_kl_loss / report_num_examples, mi,
report_rec_loss / report_num_examples, au, time.time() - start))
else:
print('epoch: %d, iter: %d, avg_loss: %.4f, kl: %.4f, recon: %.4f,' \
'time elapsed %.2fs' %
(epoch, iter_, train_loss, report_kl_loss / report_num_examples,
report_rec_loss / report_num_examples, time.time() - start))
sys.stdout.flush()
report_rec_loss = report_kl_loss = 0
report_num_examples = 0
iter_ += 1
if aggressive_flag and (iter_ % len(train_loader)) == 0:
vae.eval()
cur_mi = calc_mi(vae, val_loader)
vae.train()
if cur_mi - best_mi < 0:
mi_not_improved += 1
if mi_not_improved == 5:
aggressive_flag = False
print("STOP BURNING")
else:
best_mi = cur_mi
pre_mi = cur_mi
print('kl weight %.4f' % kl_weight)
print('epoch: %d, VAL' % epoch)
vae.eval()
with torch.no_grad():
loss, nll, kl = test(vae, val_loader, "VAL", args)
au, au_var = calc_au(vae, val_loader)
print("%d active units" % au)
# print(au_var)
if loss < best_loss:
print('update best loss')
best_loss = loss
best_nll = nll
best_kl = kl
torch.save(vae.state_dict(), args.save_path)
if loss > best_loss:
opt_dict["not_improved"] += 1
if opt_dict["not_improved"] >= decay_epoch:
opt_dict["best_loss"] = loss
opt_dict["not_improved"] = 0
opt_dict["lr"] = opt_dict["lr"] * lr_decay
vae.load_state_dict(torch.load(args.save_path))
decay_cnt += 1
print('new lr: %f' % opt_dict["lr"])
enc_optimizer = optim.Adam(vae.encoder.parameters(),
lr=opt_dict["lr"])
dec_optimizer = optim.Adam(vae.decoder.parameters(),
lr=opt_dict["lr"])
else:
opt_dict["not_improved"] = 0
opt_dict["best_loss"] = loss
if decay_cnt == max_decay:
break
if epoch % args.test_nepoch == 0:
with torch.no_grad():
loss, nll, kl = test(vae, test_loader, "TEST", args)
vae.train()
# compute importance weighted estimate of log p(x)
vae.load_state_dict(torch.load(args.save_path))
vae.eval()
with torch.no_grad():
loss, nll, kl = test(vae, test_loader, "TEST", args)
au, au_var = calc_au(vae, test_loader)
print("%d active units" % au)
# print(au_var)
test_loader = torch.utils.data.DataLoader(test_data,
batch_size=50,
shuffle=True)
with torch.no_grad():
calc_iwnll(vae, test_loader, args)
def main(args):
class uniform_initializer(object):
def __init__(self, stdv):
self.stdv = stdv
def __call__(self, tensor):
nn.init.uniform_(tensor, -self.stdv, self.stdv)
class xavier_normal_initializer(object):
def __call__(self, tensor):
nn.init.xavier_normal_(tensor)
if args.cuda:
print('using cuda')
print(args)
opt_dict = {"not_improved": 0, "lr": 1., "best_loss": 1e4}
train_data = MonoTextData(args.train_data, label=args.label)
vocab = train_data.vocab
vocab_size = len(vocab)
val_data = MonoTextData(args.val_data, label=args.label, vocab=vocab)
test_data = MonoTextData(args.test_data, label=args.label, vocab=vocab)
print('Train data: %d samples' % len(train_data))
print('finish reading datasets, vocab size is %d' % len(vocab))
print('dropped sentences: %d' % train_data.dropped)
sys.stdout.flush()
log_niter = (len(train_data) // args.batch_size) // 10
model_init = uniform_initializer(0.01)
emb_init = uniform_initializer(0.1)
device = torch.device("cuda" if args.cuda else "cpu")
args.device = device
if args.enc_type == 'lstm':
encoder = LSTMEncoder(args, vocab_size, model_init, emb_init)
args.enc_nh = args.dec_nh
else:
raise ValueError("the specified encoder type is not supported")
decoder = LSTMDecoder(args, vocab, model_init, emb_init)
vae = VAE(encoder, decoder, args).to(device)
if args.decode_from != "":
print('begin decoding')
vae.load_state_dict(torch.load(args.decode_from))
vae.eval()
save_dir = "samples/"
if not os.path.exists(save_dir):
os.makedirs(save_dir)
path = ".".join(args.decode_from.split("/")[-1].split(".")[:-1]) + \
"_{}".format(args.decoding_strategy)
with torch.no_grad():
if args.decode_input != "":
decode_data = MonoTextData(args.decode_input, vocab=vocab)
reconstruct(vae, decode_data, args.decoding_strategy,
os.path.join(save_dir, path + ".rec"), args.device)
else:
z = vae.sample_from_prior(100)
sample_from_prior(vae, z, args.decoding_strategy,
os.path.join(save_dir, path + ".sample"))
return
if args.eval:
print('begin evaluation')
vae.load_state_dict(torch.load(args.load_path))
vae.eval()
with torch.no_grad():
test_data_batch = test_data.create_data_batch(
batch_size=args.batch_size, device=device, batch_first=True)
test(vae, test_data_batch, "TEST", args)
au, au_var = calc_au(vae, test_data_batch)
print("%d active units" % au)
# print(au_var)
test_data_batch = test_data.create_data_batch(batch_size=1,
device=device,
batch_first=True)
calc_iwnll(vae, test_data_batch, args)
return
enc_optimizer = optim.SGD(vae.encoder.parameters(),
lr=1.0,
momentum=args.momentum)
dec_optimizer = optim.SGD(vae.decoder.parameters(),
lr=1.0,
momentum=args.momentum)
opt_dict['lr'] = 1.0
iter_ = decay_cnt = 0
best_loss = 1e4
best_kl = best_nll = best_ppl = 0
pre_mi = 0
aggressive_flag = True if args.aggressive else False
vae.train()
start = time.time()
kl_weight = args.kl_start
anneal_rate = (1.0 - args.kl_start) / (args.warm_up *
(len(train_data) / args.batch_size))
train_data_batch = train_data.create_data_batch(batch_size=args.batch_size,
device=device,
batch_first=True)
val_data_batch = val_data.create_data_batch(batch_size=args.batch_size,
device=device,
batch_first=True)
test_data_batch = test_data.create_data_batch(batch_size=args.batch_size,
device=device,
batch_first=True)
for epoch in range(args.epochs):
report_kl_loss = report_rec_loss = 0
report_num_words = report_num_sents = 0
for i in np.random.permutation(len(train_data_batch)):
batch_data = train_data_batch[i]
batch_size, sent_len = batch_data.size()
# not predict start symbol
report_num_words += (sent_len - 1) * batch_size
report_num_sents += batch_size
# kl_weight = 1.0
kl_weight = min(1.0, kl_weight + anneal_rate)
sub_iter = 1
batch_data_enc = batch_data
burn_num_words = 0
burn_pre_loss = 1e4
burn_cur_loss = 0
while aggressive_flag and sub_iter < 100:
enc_optimizer.zero_grad()
dec_optimizer.zero_grad()
burn_batch_size, burn_sents_len = batch_data_enc.size()
burn_num_words += (burn_sents_len - 1) * burn_batch_size
loss, loss_rc, loss_kl = vae.loss(batch_data_enc,
kl_weight,
nsamples=args.nsamples)
burn_cur_loss += loss.sum().item()
loss = loss.mean(dim=-1)
loss.backward()
torch.nn.utils.clip_grad_norm_(vae.parameters(), clip_grad)
enc_optimizer.step()
id_ = np.random.random_integers(0, len(train_data_batch) - 1)
batch_data_enc = train_data_batch[id_]
if sub_iter % 15 == 0:
burn_cur_loss = burn_cur_loss / burn_num_words
if burn_pre_loss - burn_cur_loss < 0:
break
burn_pre_loss = burn_cur_loss
burn_cur_loss = burn_num_words = 0
sub_iter += 1
# if sub_iter >= 30:
# break
# print(sub_iter)
enc_optimizer.zero_grad()
dec_optimizer.zero_grad()
loss, loss_rc, loss_kl = vae.loss(batch_data,
kl_weight,
nsamples=args.nsamples)
loss = loss.mean(dim=-1)
loss.backward()
torch.nn.utils.clip_grad_norm_(vae.parameters(), clip_grad)
loss_rc = loss_rc.sum()
loss_kl = loss_kl.sum()
if not aggressive_flag:
enc_optimizer.step()
dec_optimizer.step()
report_rec_loss += loss_rc.item()
report_kl_loss += loss_kl.item()
if iter_ % log_niter == 0:
train_loss = (report_rec_loss +
report_kl_loss) / report_num_sents
if aggressive_flag or epoch == 0:
vae.eval()
with torch.no_grad():
mi = calc_mi(vae, val_data_batch)
au, _ = calc_au(vae, val_data_batch)
vae.train()
print('epoch: %d, iter: %d, avg_loss: %.4f, kl: %.4f, mi: %.4f, recon: %.4f,' \
'au %d, time elapsed %.2fs' %
(epoch, iter_, train_loss, report_kl_loss / report_num_sents, mi,
report_rec_loss / report_num_sents, au, time.time() - start))
else:
print('epoch: %d, iter: %d, avg_loss: %.4f, kl: %.4f, recon: %.4f,' \
'time elapsed %.2fs' %
(epoch, iter_, train_loss, report_kl_loss / report_num_sents,
report_rec_loss / report_num_sents, time.time() - start))
sys.stdout.flush()
report_rec_loss = report_kl_loss = 0
report_num_words = report_num_sents = 0
iter_ += 1
if aggressive_flag and (iter_ % len(train_data_batch)) == 0:
vae.eval()
cur_mi = calc_mi(vae, val_data_batch)
vae.train()
print("pre mi:%.4f. cur mi:%.4f" % (pre_mi, cur_mi))
if cur_mi - pre_mi < 0:
aggressive_flag = False
print("STOP BURNING")
pre_mi = cur_mi
print('kl weight %.4f' % kl_weight)
vae.eval()
with torch.no_grad():
loss, nll, kl, ppl, mi = test(vae, val_data_batch, "VAL", args)
au, au_var = calc_au(vae, val_data_batch)
print("%d active units" % au)
# print(au_var)
if loss < best_loss:
print('update best loss')
best_loss = loss
best_nll = nll
best_kl = kl
best_ppl = ppl
torch.save(vae.state_dict(), args.save_path)
if loss > opt_dict["best_loss"]:
opt_dict["not_improved"] += 1
if opt_dict["not_improved"] >= decay_epoch and epoch >= 15:
opt_dict["best_loss"] = loss
opt_dict["not_improved"] = 0
opt_dict["lr"] = opt_dict["lr"] * lr_decay
vae.load_state_dict(torch.load(args.save_path))
print('new lr: %f' % opt_dict["lr"])
decay_cnt += 1
enc_optimizer = optim.SGD(vae.encoder.parameters(),
lr=opt_dict["lr"],
momentum=args.momentum)
dec_optimizer = optim.SGD(vae.decoder.parameters(),
lr=opt_dict["lr"],
momentum=args.momentum)
else:
opt_dict["not_improved"] = 0
opt_dict["best_loss"] = loss
if decay_cnt == max_decay:
break
if epoch % args.test_nepoch == 0:
with torch.no_grad():
loss, nll, kl, ppl, _ = test(vae, test_data_batch, "TEST",
args)
# save every epoch
epoch_save_pre = "%s-%03d.pt" % \
(datetime.datetime.utcnow().isoformat(), epoch)
epoch_save_fname = os.path.join("models/%s" % args.dataset,
epoch_save_pre)
torch.save(vae.state_dict(), epoch_save_fname)
vae.train()
# compute importance weighted estimate of log p(x)
vae.load_state_dict(torch.load(args.save_path))
vae.eval()
with torch.no_grad():
loss, nll, kl, ppl, _ = test(vae, test_data_batch, "TEST", args)
au, au_var = calc_au(vae, test_data_batch)
print("%d active units" % au)
# print(au_var)
test_data_batch = test_data.create_data_batch(batch_size=1,
device=device,
batch_first=True)
with torch.no_grad():
calc_iwnll(vae, test_data_batch, args)
