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 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 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 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 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 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):
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))
def load_checkpoint(path):
checkpoint = torch.load(path)
model = VAE(**checkpoint["model"])
model.load_state_dict(checkpoint["model_state_dict"])
optimizer = nn.Adam(model.parameters())
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
return model, optimizer
def load_networks(isTraining=False):
depth_net = DepthNetModel()
color_net = ColorNetModel()
d_net = VAE()
if param.useGPU:
depth_net.cuda()
color_net.cuda()
d_net.cuda()
depth_optimizer = optim.Adam(depth_net.parameters(),
lr=param.alpha,
betas=(param.beta1, param.beta2),
eps=param.eps)
color_optimizer = optim.Adam(color_net.parameters(),
lr=param.alpha,
betas=(param.beta1, param.beta2),
eps=param.eps)
d_optimizer = optim.Adam(d_net.parameters())
if isTraining:
netFolder = param.trainNet # 'TrainingData'
netName, _, _ = get_folder_content(netFolder)
net = []
for target in sorted(netName):
if target[-4:] == '.tar':
net.append(target)
if param.isContinue and net:
tokens = net[0].split('-')[1].split('.')[0]
param.startIter = int(tokens)
checkpoint = torch.load(netFolder + '/' + net[0])
depth_net.load_state_dict(checkpoint['depth_net'])
color_net.load_state_dict(checkpoint['color_net'])
d_net.load_state_dict(checkpoint['d_net'])
depth_optimizer.load_state_dict(checkpoint['depth_optimizer'])
color_optimizer.load_state_dict(checkpoint['color_optimizer'])
d_optimizer.load_state_dict(checkpoint['d_optimizer'])
else:
param.isContinue = False
else:
netFolder = param.testNet
netName, _, _ = get_folder_content(netFolder)
net = []
for target in sorted(netName):
if target[-4:] == '.tar':
net.append(target)
checkpoint = torch.load(netFolder + '/' + net[0])
depth_net.load_state_dict(checkpoint['depth_net'])
color_net.load_state_dict(checkpoint['color_net'])
d_net.load_state_dict(checkpoint['d_net'])
depth_optimizer.load_state_dict(checkpoint['depth_optimizer'])
color_optimizer.load_state_dict(checkpoint['color_optimizer'])
d_optimizer.load_state_dict(checkpoint['d_optimizer'])
return depth_net, color_net, d_net, depth_optimizer, color_optimizer, d_optimizer
def main() -> None:
tokenizer = Tokenizer(args.vocab_file)
vocabulary_size = len(tokenizer)
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()
sentence1 = input('Please input sentence1: ')
sentence2 = input('Please input sentence2: ')
s1 = [tokenizer.bos_index
] + tokenizer.encode(sentence1) + [tokenizer.eos_index]
s2 = [tokenizer.bos_index
] + tokenizer.encode(sentence2) + [tokenizer.eos_index]
z1, _ = model.encode(
torch.tensor([s1]).to(device),
torch.tensor([len(s1)]).to(device))
z2, _ = model.encode(
torch.tensor([s2]).to(device),
torch.tensor([len(s2)]).to(device))
print("\nGenerate intermediate sentences")
print(" %s" % sentence1)
for r in range(1, 10):
z = (1 - 0.1 * r) * z1 + 0.1 * r * z2
hidden = model.fc_hidden(z)
hidden = hidden.view(1, -1,
model.dim_hidden).transpose(0, 1).contiguous()
start_predictions = torch.zeros(1, 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)
tokens = predictions[0, 0]
tokens = tokens[tokens != tokenizer.eos_index].tolist()
print("[%d:%d] %s" % (10 - r, r, tokenizer.decode(tokens)))
print(" %s" % sentence2)
def main():
parser = argparse.ArgumentParser(description='Train VAE.')
parser.add_argument('-c', '--config', help='Config file.')
args = parser.parse_args()
print(args)
c = json.load(open(args.config))
print(c)
lookback = 50 # 160
input_dim = 1
returns_lookback = 20
start_date = datetime.strptime(c['start_date'],
'%Y/%m/%d') if c['start_date'] else None
end_date = datetime.strptime(c['end_date'],
'%Y/%m/%d') if c['end_date'] else None
max_n_files = None
out_path = Path(c['out_dir'])
out_path.mkdir(exist_ok=True)
# setup the VAE
vae = VAE(c['vae_series_length'], z_dim=c['z_dim'], use_cuda=c['cuda'])
if c['checkpoint_load']:
checkpoint = torch.load(c['checkpoint_load'])
vae.load_state_dict(checkpoint['model_state_dict'])
ticker_files = glob(str(Path(c['in_dir']) / '*.csv'))
if 0:
ticker_files = ticker_files[:100]
print(f'Found {len(ticker_files)} ticker files.')
extract_vae_distances(start_date=start_date,
end_date=end_date,
n_lookback_days=c['n_lookback_days'],
n_backtest_days=c['n_backtest_days'],
n_trade_days=c['n_trade_days'],
n_pairs_vae=c['n_pairs_vae'],
n_pairs_backtest=c['n_pairs_backtest'],
vae=vae,
returns_lookback=returns_lookback,
ticker_files=ticker_files,
fundamentals_file=c['fundamentals_file'],
out_dir=c['out_dir'],
r_script_exe=c['r_script_exe'],
r_backtest_script=c['r_backtest_script'],
r_trade_script=c['r_trade_script'],
backtest_sd=c['backtest_sd'],
backtest_returns_file=c['backtest_returns_file'],
backtest_plot_file=c['backtest_plot_file'],
trade_returns_file=c['trade_returns_file'],
trade_plot_file=c['trade_plot_file'],
cuda=c['cuda'])
def load_checkpoint(file_path, use_cuda=False):
"""Return EmbedNet instance"""
if use_cuda:
checkpoint = torch.load(file_path)
else:
checkpoint = torch.load(file_path,
map_location=lambda storage, location: storage)
n_latents = checkpoint['n_latents']
vae = VAE(n_latents=n_latents)
vae.load_state_dict(checkpoint['state_dict'])
return vae
def show_reconstruction():
vae = VAE()
vae.load_state_dict(torch.load("./vae.pth"))
it = iter(test_dataloader)
imgs, _ = it.next()
imgs1 = torchvision.utils.make_grid(imgs)
plt.subplot(1, 2, 1)
showimg(imgs1)
out, _, _ = vae(imgs)
out = torchvision.utils.make_grid(out)
plt.subplot(1, 2, 2)
showimg(out.detach())
plt.show()
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 slave(comm):
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'])
controller = Controller()
controller.load_state_dict(
torch.load(cfg.ctrl_save_ckpt,
map_location=lambda storage, loc: storage)['model'])
env = DoomTakeCover(False)
rewards = []
for epi in range(cfg.trials_per_pop * 4):
obs = env.reset()
model.reset()
for step in range(cfg.max_steps):
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
rewards.append(step)
print('Workder {} got reward {} at epi {}'.format(
comm.rank, step, epi))
rewards = np.array(rewards)
comm.send(rewards, dest=0, tag=1)
print('Worker {} sent rewards to master'.format(comm.rank))
def test_vae():
data = glob.glob('../../data/doom_frames/*.npz')
data = np.random.choice(data)
frames = np.load(data)['sx']
model = VAE().cuda(3)
stat_dict = torch.load(cfg.vae_save_ckpt)['model']
model.load_state_dict(stat_dict)
x = frames.transpose(0, 3, 1, 2)
x = torch.from_numpy(x).float().cuda(3) / 255.0
_, _, x_rec, _ = model(x)
x_rec = x_rec.detach().cpu().numpy().transpose(0, 2, 3, 1) * 255.0
new_frames = np.zeros((x_rec.shape[0], x_rec.shape[1] * 2 + 20,
x_rec.shape[2], x_rec.shape[3]))
new_frames[:, :x_rec.shape[1], :, :] = x_rec
new_frames[:, -x_rec.shape[1]:, :, :] = frames.astype(np.float)
print(new_frames.shape, 'vae.avi')
write_video(new_frames, 'vae.avi', (64, 148))
os.system('mv vae.avi /home/bzhou/Dropbox/share')
def main(config):
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
model = VAE().cuda()
state = torch.load(config.ckp_path)
model.load_state_dict(state['state_dict'])
os.makedirs(os.path.dirname(config.save_path), exist_ok=True)
torch.manual_seed(66666)
np.random.seed(66666)
random.seed(66666)
z = torch.randn((32, 512)).cuda()
predict = model.decode(z)
torchvision.utils.save_image(predict.data,
config.save_path,
nrow=8,
normalize=True)
class Generator(object):
def __init__(self):
_, _, self.vocab = get_iterators(opt)
self.vae = VAE(opt)
self.vae.embedding.weight.data.copy_(self.vocab.vectors)
self.vae = get_cuda(self.vae)
checkpoint = T.load('data/saved_models/vae_model.121.pyt')
self.vae.load_state_dict(checkpoint['vae_dict'])
self.vae.eval()
del checkpoint
def generate(self, encodings):
sentences = []
for z in encodings.numpy():
z = get_cuda(T.from_numpy(z)).view((1, -1))
h_0 = get_cuda(T.zeros(opt.n_layers_G, 1, opt.n_hidden_G))
c_0 = get_cuda(T.zeros(opt.n_layers_G, 1, opt.n_hidden_G))
G_hidden = (h_0, c_0)
G_inp = T.LongTensor(1, 1).fill_(self.vocab.stoi[opt.start_token])
G_inp = get_cuda(G_inp)
sentence = opt.start_token + " "
num_words = 0
while G_inp[0][0].item() != self.vocab.stoi[opt.end_token]:
with T.autograd.no_grad():
logit, G_hidden, _ = self.vae(None, G_inp, z, G_hidden)
probs = F.softmax(logit[0] / TEMPERATURE, dim=1)
G_inp = T.multinomial(probs, 1)
sentence += (self.vocab.itos[G_inp[0][0].item()] + " ")
num_words += 1
if num_words > 64:
break
sentence = sentence.replace('<unk>',
'').replace('<sos>', '').replace(
'<eos>', '').replace('<pad>', '')
sentences.append(sentence)
return sentences
def main() -> None:
tokenizer = Tokenizer(args.vocab_file)
vocabulary_size = len(tokenizer)
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()
z = torch.randn(args.sample_size, args.dim_latent, device=device)
hidden = model.fc_hidden(z)
hidden = hidden.view(args.sample_size, -1,
model.dim_hidden).transpose(0, 1).contiguous()
start_predictions = torch.zeros(args.sample_size, 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 pred in predictions:
tokens = pred[0]
tokens = tokens[tokens != tokenizer.eos_index].tolist()
print(tokenizer.decode(tokens))
clusterIndexList.append([])
for i in range(len(listResult)):
clusterIndexList[listResult[i]].append(i)
reconNew = np.zeros(
(scData.features.shape[0], scData.features.shape[1]))
# Convert to Tensor
reconNew = torch.from_numpy(reconNew)
if args.precisionModel == 'Double':
reconNew = reconNew.type(torch.DoubleTensor)
elif args.precisionModel == 'Float':
reconNew = reconNew.type(torch.FloatTensor)
reconNew = reconNew.to(device)
model.load_state_dict(ptstatus)
for clusterIndex in clusterIndexList:
reconUsage = recon[clusterIndex]
scDataInter = scDatasetInter(reconUsage)
train_loader = DataLoader(
scDataInter, batch_size=args.batch_size, shuffle=False, **kwargs)
for epoch in range(1, args.cluster_epochs + 1):
reconCluster, originalCluster, zCluster = train(
epoch, EMFlag=True)
count = 0
for i in clusterIndex:
reconNew[i] = reconCluster[count, :]
count += 1
# empty cuda cache
del originalCluster
dropout_rate=0.5,
**learning_kwargs)
model.update_prior()
run(opts=[optimizer_decoder],
n_epochs=args.n_dec_epochs,
dropout_rate=0,
**learning_kwargs)
train_scores.append(
evaluate(model, train_data, train_data, metrics, 0.01)[0])
valid_scores.append(
evaluate(model, valid_in_data, valid_out_data, metrics, 1)[0])
if valid_scores[-1] > best_ndcg:
best_ndcg = valid_scores[-1]
model_best.load_state_dict(deepcopy(model.state_dict()))
print(
f'epoch {epoch} | valid ndcg@100: {valid_scores[-1]:.4f} | ' +
f'best valid: {best_ndcg:.4f} | train ndcg@100: {train_scores[-1]:.4f}'
)
test_metrics = [NDCGK(100), RecallK(20), RecallK(50)]
final_scores = evaluate_recpack(model_best, test_in_data, test_out_data,
test_metrics)
print("Original results")
for metric, score in zip(test_metrics, final_scores):
print(f"{metric.name}:\t{score:.4f}")
def test():
parser = argparse.ArgumentParser(description='Train VAE.')
parser.add_argument('-c', '--config', help='Config file.')
args = parser.parse_args()
print(args)
c = json.load(open(args.config))
print(c)
# clear param store
pyro.clear_param_store()
# batch_size = 64
# root_dir = r'D:\projects\trading\mlbootcamp\tickers'
# series_length = 60
lookback = 50 # 160
input_dim = 1
test_start_date = datetime.strptime(
c['test_start_date'], '%Y/%m/%d') if c['test_start_date'] else None
test_end_date = datetime.strptime(
c['test_end_date'], '%Y/%m/%d') if c['test_end_date'] else None
min_sequence_length_test = 2 * (c['series_length'] + lookback)
max_n_files = None
out_path = Path(c['out_dir'])
out_path.mkdir(exist_ok=True)
# load_path = 'out_saved/checkpoint_0035.pt'
dataset_test = create_ticker_dataset(
c['in_dir'],
c['series_length'],
lookback,
min_sequence_length_test,
start_date=test_start_date,
end_date=test_end_date,
fixed_start_date=True,
normalised_returns=c['normalised_returns'],
max_n_files=max_n_files)
test_loader = DataLoader(dataset_test,
batch_size=c['batch_size'],
shuffle=False,
num_workers=0,
drop_last=True)
# N_train_data = len(dataset_train)
N_test_data = len(dataset_test)
# N_mini_batches = N_train_data // c['batch_size']
# N_train_time_slices = c['batch_size'] * N_mini_batches
print(f'N_test_data: {N_test_data}')
# setup the VAE
vae = VAE(c['series_length'], z_dim=c['z_dim'], use_cuda=c['cuda'])
# setup the optimizer
# adam_args = {"lr": args.learning_rate}
# optimizer = Adam(adam_args)
# setup the inference algorithm
# elbo = JitTrace_ELBO() if args.jit else Trace_ELBO()
# svi = SVI(vae.model, vae.guide, optimizer, loss=elbo)
if c['checkpoint_load']:
checkpoint = torch.load(c['checkpoint_load'])
vae.load_state_dict(checkpoint['model_state_dict'])
# optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
if 1:
find_similar(vae, test_loader, c['cuda'])
# Visualise first batch.
batch = next(iter(test_loader))
x = batch['series']
if c['cuda']:
x = x.cuda()
x = x.float()
x_reconst = vae.reconstruct_img(x)
x = x.cpu().numpy()
x_reconst = x_reconst.cpu().detach().numpy()
n = min(5, x.shape[0])
fig, axes = plt.subplots(n, 1, squeeze=False)
for s in range(n):
ax = axes[s, 0]
ax.plot(x[s])
ax.plot(x_reconst[s])
fig.savefig(out_path / f'test_batch.png')
img_fn = str(int(time.time())) + ".png"
plt.imsave(os.path.join(output_fp, img_fn), img)
print("|Picture is generated. |{}".format(img_fn))
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument("-m",
"--model",
type=str,
required=True,
help="model file path")
parser.add_argument("-o",
"--output",
type=str,
required=True,
help="output file path")
model_fp = parser.parse_args().model
output_fp = parser.parse_args().output
vae = VAE()
vae.load_state_dict(tor.load(model_fp))
rand_v = tor.randn((1, 512))
#rand_v = tor.FloatTensor(1, 512).uniform_(0,1)
random_generator(vae, rand_v, output_fp)
return options.parse_args()
args = setup_args()
if not torch.cuda.is_available():
args.use_gpu = False
os.makedirs(args.save_dir, exist_ok=True)
# ============= TRAINING INITIALIZATION ==============
# initialize autoencoder
netRNA = FC_VAE(n_input=7633, nz=args.latent_dims)
netImage = VAE(latent_variable_size=args.latent_dims, batchnorm=True)
netImage.load_state_dict(torch.load(args.pretrained_file))
print("Pre-trained model loaded from %s" % args.pretrained_file)
if args.conditional_adv:
netClf = FC_Classifier(nz=args.latent_dims + 10)
assert not args.conditional
else:
netClf = FC_Classifier(nz=args.latent_dims)
if args.conditional:
netCondClf = Simple_Classifier(nz=args.latent_dims)
if args.use_gpu:
netRNA.cuda()
netImage.cuda()
netClf.cuda()
for i, (img, target) in enumerate(tqdm(dataloader)):
output, _, _ = model.encoder(img)
output = torch.squeeze(output).detach().numpy()
plt.scatter(output[:,0], output[:,1], color=colors[target], s=4)
plt.show()
def random_sample(model, cmap='gray'):
n_to_show = 20
znew = torch.randn((n_to_show, 2))
recon = model.decoder(znew)
for i in range(n_to_show):
plt.subplot(4, 5, i+1)
plt.imshow(torch.squeeze(recon[i]).detach().numpy(), cmap)
plt.xticks([]); plt.yticks([])
plt.show()
if __name__=='__main__':
test_dataset = datasets.MNIST(root='./mnist_data/', train=False, transform=transforms.ToTensor(), download=False)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=256, shuffle=False)
model = VAE()
model.load_state_dict(torch.load('./checkpoints/epoch_050.pth'))
show_2D_feature(model, test_loader)
random_sample(model)
# show_imgs(model, test_loader)
class TrainingModel(object):
def __init__(self, args, config):
self.__dict__.update(config)
self.config = config
random.seed(self.seed)
torch.manual_seed(self.seed)
np.random.seed(self.seed)
if use_cuda:
torch.cuda.manual_seed(self.seed)
torch.cuda.manual_seed_all(self.seed)
torch.cuda.set_device(args.gpu)
#torch.backends.cudnn.benchmark = False
#torch.backends.cudnn.deterministic = True
self.message = args.m
self.data_generator = DataGenerator(self.config)
self.vocab_size = self.data_generator.vocab_size
self.ent_size = self.data_generator.ent_size
self.model_name = 'IERM'
if args.m != "":
self.saveModeladdr = './trainModel/checkpoint_%s_%s.pkl' % (
self.model_name, args.m)
else:
self.saveModeladdr = './trainModel/' + args.save
self.model = Ranker(self.vocab_size, self.ent_size, self.config)
self.VAE_model = VAE(self.vocab_size, self.ent_size,
self.model.word_emb, self.model.ent_emb,
self.config)
if use_cuda:
self.model.cuda()
self.VAE_model.cuda()
vae_lr = self.config[
'pretrain_lr'] if config['pretrain_step'] > 0 else config['vae_lr']
self.vae_optimizer = getOptimizer(config['vae_optim'],
self.VAE_model.parameters(),
lr=vae_lr,
betas=(0.99, 0.99))
self.ranker_optimizer = getOptimizer(
config['ranker_optim'],
self.model.parameters(),
lr=config['ranker_lr'],
weight_decay=config['weight_decay'])
vae_model_size = sum(p.numel() for p in self.VAE_model.parameters())
ranker_size = sum(p.numel() for p in self.model.parameters())
#print 'Model size: ', vae_model_size, ranker_size
#exit(-1)
if args.resume and os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
#print checkpoint.keys()
self.model.load_state_dict(checkpoint['rank_state_dict'])
self.VAE_model.load_state_dict(checkpoint['vae_state_dict'])
self.vae_optimizer.load_state_dict(checkpoint['vae_optimizer'])
self.ranker_optimizer.load_state_dict(checkpoint['rank_optimizer'])
else:
print("Creating a new model")
self.timings = defaultdict(list) #record the loss iterations
self.evaluator = rank_eval()
self.epoch = 0
self.step = 0
self.kl_weight = 1
if args.visual:
self.config['visual'] = True
self.writer = SummaryWriter('runs/' + args.m)
else:
self.config['visual'] = False
self.reconstr_loss = nn.MSELoss()
def add_values(self, iter, value_dict):
for key in value_dict:
self.writer.add_scalar(key, value_dict[key], iter)
def adjust_learning_rate(self, optimizer, lr, decay_rate=.5):
for param_group in optimizer.param_groups:
param_group['lr'] = lr * decay_rate
def kl_anneal_function(self, anneal_function, step, k=0.0025, x0=2500):
if anneal_function == 'logistic':
return float(1 / (1 + np.exp(-k * (step - x0))))
elif anneal_function == 'linear':
return min(1, step / x0)
def vae_loss(self, input_qw, reconstr_w, input_qe, reconstr_e, prior_mean,
prior_var, posterior_mean, posterior_var, posterior_log_var):
# Reconstruction term
if self.config['reconstruct'] != 'entity':
input_qw_bow = to_bow(input_qw, self.vocab_size)
input_qw_bow = Tensor2Varible(torch.tensor(input_qw_bow).float())
#reconstr_w = torch.log_softmax(reconstr_w + 1e-10,dim=1)
#RL_w = -torch.sum(input_qw_bow * reconstr_w , dim=1)
#RL_w = self.reconstr_loss(reconstr_w,input_qw_bow)
RL_w = -torch.sum(
input_qw_bow * reconstr_w +
(1 - input_qw_bow) * torch.log(1 - torch.exp(reconstr_w)),
dim=1)
else:
RL_w = Tensor2Varible(torch.tensor([0]).float())
if self.config['reconstruct'] != 'word':
input_qe_bow = to_bow(input_qe, self.ent_size)
input_qe_bow = Tensor2Varible(torch.tensor(input_qe_bow).float())
#RL_e = -torch.sum(input_qe_bow * reconstr_e, dim=1)
#RL_e = self.reconstr_loss(reconstr_e,input_qe_bow)
RL_e = -torch.sum(
input_qe_bow * reconstr_e +
(1 - input_qe_bow) * torch.log(1 - torch.exp(reconstr_e)),
dim=1)
else:
RL_e = Tensor2Varible(torch.tensor([0]).float())
# KL term
# var division term
var_division = torch.sum(posterior_var / prior_var, dim=1)
# diff means term
diff_means = prior_mean - posterior_mean
diff_term = torch.sum((diff_means * diff_means) / prior_var, dim=1)
# logvar det division term
logvar_det_division = \
prior_var.log().sum() - posterior_log_var.sum(dim=1)
# combine terms
KL = 0.5 * (var_division + diff_term - self.model.intent_num +
logvar_det_division)
loss = self.kl_weight * KL + RL_w + RL_e
#loss = 0.001 * KL + RL_w + RL_e
return loss.sum(), KL.sum(), RL_w.sum(), RL_e.sum()
def pretraining(self):
if self.pretrain_step <= 0:
return
train_start_time = time.time()
data_reader = self.data_generator.pretrain_reader(self.pretrain_bs)
total_loss = 0.
total_KL_loss = 0.
total_RLw_loss = 0.
total_RLe_loss = 0.
for step in xrange(self.pretrain_step):
input_qw, input_qe = next(data_reader)
#self.kl_weight = self.kl_anneal_function('logistic', step)
topic_e, vae_loss, kl_loss, rl_w_loss, rl_e_loss = self.train_VAE(
input_qw, input_qe)
vae_loss.backward()
torch.nn.utils.clip_grad_value_(
self.VAE_model.parameters(),
self.clip_grad) # clip_grad_norm(, )
self.vae_optimizer.step()
vae_loss = vae_loss.data
#print ('VAE loss: %.3f\tKL: %.3f\tRL_w:%.3f\tRL_e:%.3f' % (vae_loss, kl_loss, rl_w_loss, rl_e_loss))
if torch.isnan(vae_loss):
print("Got NaN cost .. skipping")
exit(-1)
continue
#if self.config['visual']:
# self.add_values(step, {'vae_loss': vae_loss, 'kl_loss': kl_loss, 'rl_w_loss': rl_w_loss,
# 'rl_e_loss': rl_e_loss, 'kl_weight': self.kl_weight})
total_loss += vae_loss
total_KL_loss += kl_loss
total_RLw_loss += rl_w_loss
total_RLe_loss += rl_e_loss
if step != 0 and step % self.pretrain_freq == 0:
total_loss /= self.pretrain_freq
total_KL_loss /= self.pretrain_freq
total_RLw_loss /= self.pretrain_freq
total_RLe_loss /= self.pretrain_freq
print('Step: %d\t Elapsed:%.2f' %
(step, time.time() - train_start_time))
print(
'Pretrain VAE loss: %.3f\tKL: %.3f\tRL_w:%.3f\tRL_e:%.3f' %
(total_loss, total_KL_loss, total_RLw_loss,
total_RLe_loss))
if self.config['visual']:
self.add_values(
step, {
'vae_loss': total_loss,
'kl_loss': total_KL_loss,
'rl_w_loss': total_RLw_loss,
'rl_e_loss': total_RLe_loss,
'kl_weight': self.kl_weight
})
total_loss = 0.
total_KL_loss = 0.
total_RLw_loss = 0.
total_RLe_loss = 0.
print '=============================================='
#self.generate_beta_phi_3(show_topic_limit=5)
self.save_checkpoint(message=self.message + '-pretraining')
print('Pretraining end')
#recovering the learning rate
self.adjust_learning_rate(self.vae_optimizer, self.config['vae_lr'], 1)
def trainIters(self, ):
self.step = 0
train_start_time = time.time()
patience = self.patience
best_ndcg10 = 0.0
last_ndcg10 = 0.0
data_reader = self.data_generator.pair_reader(self.batch_size)
total_loss = 0.0
total_rank_loss = 0.
total_vae_loss = 0.
total_KL_loss = 0.
total_RLw_loss = 0.
total_RLe_loss = 0.
for step in xrange(self.steps):
out = next(data_reader)
input_qw, input_qe, input_dw_pos, input_de_pos, input_dw_neg, input_de_neg = out
rank_loss, vae_total_loss, KL_loss, RL_w_loss, RL_e_loss \
= self.train(input_qw,input_qe,input_dw_pos,input_de_pos,input_dw_neg,input_de_neg)
cur_total_loss = rank_loss + vae_total_loss
if torch.isnan(cur_total_loss):
print("Got NaN cost .. skipping")
continue
self.step += 1
total_loss += cur_total_loss
total_rank_loss += rank_loss
total_vae_loss += vae_total_loss
total_KL_loss += KL_loss
total_RLw_loss += RL_w_loss
total_RLe_loss += RL_e_loss
if self.eval_freq != -1 and self.step % self.eval_freq == 0:
with torch.no_grad():
valid_performance = self.test(
valid_or_test='valid',
source=self.config['click_model'])
current_ndcg10 = valid_performance['ndcg@10']
if current_ndcg10 > best_ndcg10:
print 'Got better result, save to %s' % self.saveModeladdr
best_ndcg10 = current_ndcg10
patience = self.patience
self.save_checkpoint(message=self.message)
#self.generate_beta_phi_3(show_topic_limit=5)
elif current_ndcg10 <= last_ndcg10 * self.cost_threshold:
patience -= 1
last_ndcg10 = current_ndcg10
if self.step % self.train_freq == 0:
total_loss /= self.train_freq
total_rank_loss /= self.train_freq
total_vae_loss /= self.train_freq
total_KL_loss /= self.train_freq
total_RLw_loss /= self.train_freq
total_RLe_loss /= self.train_freq
self.timings['train'].append(total_loss)
print('Step: %d\t Elapsed:%.2f' %
(step, time.time() - train_start_time))
print(
'Train total loss: %.3f\tRank loss: %.3f\tVAE loss: %.3f' %
(total_loss, total_rank_loss, total_vae_loss))
print('KL loss: %.3f\tRL W: %.3f\tRL E: %.3f' %
(total_KL_loss, total_RLw_loss, total_RLe_loss))
print('Patience left: %d' % patience)
if self.config['visual']:
self.add_values(
step, {
'Train vae_loss': total_loss,
'Train kl_loss': total_KL_loss,
'Train rl_w_loss': total_RLw_loss,
'Train rl_e_loss': total_RLe_loss,
'Train Rank loss': total_rank_loss
})
total_loss = 0
total_rank_loss = 0.
total_vae_loss = 0.
total_KL_loss = 0.
total_RLw_loss = 0.
total_RLe_loss = 0.
if patience < 0:
print 'patience runs out...'
break
print 'Patience___: ', patience
print("All done, exiting...")
def test(self, valid_or_test, source):
predicted = []
results = defaultdict(list)
if valid_or_test == 'valid':
is_test = False
data_addr = self.valid_rank_addr
data_source = self.data_generator.pointwise_reader_evaluation(
data_addr, is_test=is_test, label_type=source)
elif valid_or_test == 'ntcir13' or valid_or_test == 'ntcir14':
is_test = True
data_source = self.data_generator.pointwise_ntcir_generator(
valid_or_test)
source = 'HUMAN'
else:
is_test = True
data_addr = self.test_rank_addr
data_source = self.data_generator.pointwise_reader_evaluation(
data_addr, is_test=is_test, label_type=source)
start = time.clock()
count = 0
for out in data_source:
(qid, dids, input_qw, input_qe, input_dw, input_de, gt_rels) = out
gt_rels = map(lambda t: score2cutoff(source, t), gt_rels)
rels_predicted = self.predict(input_qw, input_qe, input_dw,
input_de).view(-1).cpu().numpy()
result = self.evaluator.eval(gt_rels, rels_predicted)
for did, gt, pred in zip(dids, gt_rels, rels_predicted):
predicted.append((qid, did, pred, gt))
for k, v in result.items():
results[k].append(v)
count += 1
elapsed = (time.clock() - start)
print('Elapsed:%.3f\tAvg:%.3f' % (elapsed, elapsed / count))
performances = {}
for k, v in results.items():
performances[k] = np.mean(v)
print '------Source: %s\tPerformance-------:' % source
print 'Validating...' if valid_or_test == 'valid' else 'Testing'
print 'Message: %s' % self.message
print 'Source: %s' % source
print performances
if valid_or_test != 'valid':
path = './results/' + self.message + '_' + valid_or_test + '_' + source
if not os.path.exists(path):
os.makedirs(path)
out_file = open('%s/%s.predicted.txt' % (path, self.model_name),
'w')
for qid, did, pred, gt in predicted:
print >> out_file, '\t'.join([qid, did, str(pred), str(gt)])
return performances
def get_text(self, input, map_fun):
text_list = []
for element in input:
if element == 0:
break
text_list.append(map_fun(element))
return ' '.join(text_list)
def generate_beta_phi_3(self, topK=10, show_topic_limit=-1):
beta, phi = self.VAE_model.infer_topic_dis(topK)
topics = defaultdict(list)
topics_ents = defaultdict(list)
show_topic_num = self.config[
'intent_num'] if show_topic_limit == -1 else show_topic_limit
for i in range(show_topic_num):
idxs = beta[i]
eidxs = phi[i]
component_words = [
self.data_generator.id2word[idx] for idx in idxs.cpu().numpy()
]
component_ents = [
self.data_generator.id2ent[self.data_generator.new2old[idx]]
for idx in eidxs.cpu().numpy()
]
topics[i] = component_words
topics_ents[i] = component_ents
print '--------Topic-Word-------'
prefix = ('./topic/%s/' % args.m)
if not os.path.exists(prefix):
os.makedirs(prefix)
outfile = open(prefix + 'topic-words.txt', 'w')
for k in topics:
print >> outfile, (str(k) + ' : ' + ' '.join(topics[k]))
print >> outfile, (str(k) + ' : ' + ' '.join(topics_ents[k]))
return topics, topics_ents
def run_test_topic(self, out_file_name, topK, topicNum):
topics_words, topics_ents = self.generate_beta_phi_3(topK)
data_addr = self.test_rank_addr
data_source = self.data_generator.pointwise_reader_evaluation(
data_addr, is_test=True, label_type=self.config['click_model'])
out_file = open(out_file_name, 'w')
with torch.no_grad():
self.VAE_model.eval()
self.model.eval()
for i, out in enumerate(data_source):
(qid, dids, input_qw, input_qe, input_dw, input_de,
gt_rels) = out
theta = self.VAE_model.get_theta(input_qw, input_qe)
input_qw = input_qw[0]
input_qe = input_qe[0]
input_w = self.get_text(
input_qw, lambda w: self.data_generator.id2word[w])
input_e = self.get_text(
input_qe, lambda e: self.data_generator.id2ent[
self.data_generator.new2old[e]])
theta = theta[0].data.cpu().numpy()
top_indices = np.argsort(theta)[::-1][:3]
#print '========================='
print >> out_file, 'Query: ', input_w
print >> out_file, 'Entity: ', input_e
for j, k in enumerate(top_indices):
ws = topics_words[k]
es = topics_ents[k]
print >> out_file, '%d Word Topic %d: %s' % (j, k,
' '.join(ws))
print >> out_file, '%d Entity Topic %d: %s' % (
j, k, ' '.join(es))
def generate_topic_word_ent(self, out_file, topK=10):
print 'Visualizing ...'
data_addr = self.test_rank_addr
data_source = self.data_generator.pointwise_reader_evaluation(
data_addr, is_test=True, label_type=self.config['click_model'])
out_file = open(out_file, 'w')
with torch.no_grad():
self.VAE_model.eval()
self.model.eval()
for i, out in enumerate(data_source):
(input_qw, input_qe, input_dw, input_de, gt_rels) = out
_, word_indices, ent_indices = self.VAE_model.get_topic_words(
input_qw, input_qe, topK=topK)
word_indices = word_indices[0].data.cpu().numpy()
ent_indices = ent_indices[0].data.cpu().numpy()
#print 'ent_indices: ', ent_indices
#print 'word_indices: ', word_indices
input_qw = input_qw[0]
input_qe = input_qe[0]
input_w = self.get_text(
input_qw, lambda w: self.data_generator.id2word[w])
input_e = self.get_text(
input_qe, lambda e: self.data_generator.id2ent[
self.data_generator.new2old[e]])
reconstuct_w = self.get_text(
word_indices, lambda w: self.data_generator.id2word[w])
reconstuct_e = self.get_text(
ent_indices, lambda e: self.data_generator.id2ent[
self.data_generator.new2old[e]])
print >> out_file, ('%d: Word: %s\tRecons: %s' %
(i + 1, input_w, reconstuct_w))
print >> out_file, ('%d: Ent: %s\tRecons: %s' %
(i + 1, input_e, reconstuct_e))
def train_VAE(self, input_qw, input_qe):
self.VAE_model.train()
self.VAE_model.zero_grad()
self.vae_optimizer.zero_grad()
topic_embeddings, logPw, logPe, prior_mean, prior_variance,\
poster_mu, poster_sigma, poster_log_sigma = self.VAE_model(input_qw,input_qe)
vae_total_loss, KL, RL_w, RL_e = self.vae_loss(
input_qw, logPw, input_qe, logPe, prior_mean, prior_variance,
poster_mu, poster_sigma, poster_log_sigma)
#vae_total_loss.backward(retain_graph=True)
# `clip_grad_norm` helps prevent the exploding gradient problem in RNNs / LSTMs.
#torch.nn.utils.clip_grad_value_(self.VAE_model.parameters(), self.clip_grad) # clip_grad_norm(, )
#self.vae_optimizer.step()
return topic_embeddings, vae_total_loss, KL.data, RL_w.data, RL_e.data
def train(self, input_qw, input_qe, input_dw_pos, input_de_pos,
input_dw_neg, input_de_neg):
# Turn on training mode which enables dropout.
self.model.train()
self.model.zero_grad()
self.ranker_optimizer.zero_grad()
topic_embeddings, vae_total_loss, KL_loss, RL_w_loss, RL_e_loss = self.train_VAE(
input_qw, input_qe)
score_pos, orth_loss_1 = self.model(input_qw, input_qe, input_dw_pos,
input_de_pos, topic_embeddings)
score_neg, orth_loss_2 = self.model(input_qw, input_qe, input_dw_neg,
input_de_neg, topic_embeddings)
rank_loss = torch.sum(torch.clamp(1.0 - score_pos + score_neg, min=0))
vae_weight = self.config['intent_lambda']
orth_loss = (orth_loss_1 + orth_loss_2) / 2
total_loss = rank_loss + vae_weight * vae_total_loss + orth_loss
total_loss.backward()
## update parameters
# `clip_grad_norm` helps prevent the exploding gradient problem in RNNs / LSTMs.
torch.nn.utils.clip_grad_value_(self.VAE_model.parameters(),
self.clip_grad) # clip_grad_norm(, )
torch.nn.utils.clip_grad_value_(self.model.parameters(),
self.clip_grad) #clip_grad_norm(, )
self.ranker_optimizer.step()
self.vae_optimizer.step()
return rank_loss.data, vae_total_loss.data, KL_loss, RL_w_loss, RL_e_loss
def predict(self, input_qw, input_qe, input_dw, input_de):
# Turn on evaluation mode which disables dropout.
with torch.no_grad():
self.VAE_model.eval()
self.model.eval()
topic_embeddings = self.VAE_model(input_qw, input_qe)
rels_predicted, _ = self.model(input_qw, input_qe, input_dw,
input_de, topic_embeddings)
return rels_predicted
def save_checkpoint(self, message):
filePath = os.path.join(self.saveModeladdr)
#if not os.path.exists(filePath):
# os.makedirs(filePath)
torch.save(
{
'vae_state_dict': self.VAE_model.state_dict(),
'rank_state_dict': self.model.state_dict(),
'vae_optimizer': self.vae_optimizer.state_dict(),
'rank_optimizer': self.ranker_optimizer.state_dict()
}, filePath)
def get_embeddings(self):
word_embeddings = self.model.word_emb.weight.detach().cpu().numpy()
ent_embeddings = self.model.ent_emb.weight.detach().cpu().numpy()
topic_embeddings = self.model.topic_embedding.detach().cpu().numpy()
print 'Topic size: ', topic_embeddings.shape[0]
cPickle.dump((word_embeddings, ent_embeddings, topic_embeddings),
open('./topic_analysis/w_e_t_embedding.pkl', 'w'))
print 'saved'
return
class Trainer(object):
def __init__(self, train_loader, test_loader, config):
self.train_loader = train_loader
self.test_loader = test_loader
self.config = config
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
self.num_epochs = config.num_epochs
self.lr = config.lr
self.in_channel = config.in_channel
self.image_size = config.image_size
self.hidden_dim = config.hidden_dim
self.output_dim = config.output_dim
self.log_interval = config.log_interval
self.sample_interval = config.sample_interval
self.ckpt_interval = config.ckpt_interval
self.sample_folder = config.sample_folder
self.ckpt_folder = config.ckpt_folder
self.build_net()
self.vis = Visualizer()
def build_net(self):
# define network
self.net = VAE(self.in_channel, self.image_size, self.hidden_dim,
self.output_dim)
if self.config.mode == 'test' and self.config.training_path == '':
print("[*] Enter model path!")
exit()
# if training model exists
if self.config.training_path != '':
self.net.load_state_dict(
torch.load(self.config.training_path,
map_location=lambda storage, loc: storage))
print("[*] Load weight from {}!".format(self.config.training_path))
self.net.to(self.device)
# define loss function
def loss_function(self, recon_x, x, mu, logvar):
criterion = nn.MSELoss(reduction='sum').to(self.device)
bce = criterion(recon_x,
x.view(-1, self.in_channel * (self.image_size**2)))
kld = -0.5 * torch.sum(1 + logvar - mu**2 - logvar.exp())
return bce + kld
def train(self):
# define optimizer
optimizer = Adam(self.net.parameters(), self.lr)
step = 0
print("[*] Learning started!")
# get fixed sample
temp_iter = iter(self.train_loader)
fixed_imgs, _ = next(temp_iter)
fixed_imgs = fixed_imgs.to(self.device)
# save fixed sample image
x_path = os.path.join(self.sample_folder, 'fixed_input.png')
save_image(fixed_imgs, x_path, normalize=True)
print("[*] Save fixed input image!")
# flatten
fixed_imgs = fixed_imgs.view(fixed_imgs.size(0), -1)
for epoch in range(self.num_epochs):
for i, (imgs, _) in enumerate(self.train_loader):
self.net.train()
imgs = imgs.view(imgs.size(0), -1)
imgs = imgs.to(self.device)
# forwarding and compute loss
recon, mu, logvar = self.net(imgs)
# testing code
# print("reconstructed:", recon.shape)
# print("original:", imgs.shape)
loss = self.loss_function(recon, imgs, mu, logvar)
# backwarding
optimizer.zero_grad()
loss.backward()
optimizer.step()
# do logging
if (step + 1) % self.log_interval == 0:
print("[{}/{}] [{}/{}] Loss:{:3f}".format(
epoch + 1, self.num_epochs, i + 1,
len(self.train_loader),
loss.item() / len(imgs)))
self.vis.plot("loss plot", loss.item() / len(imgs))
# do sampling
if (step + 1) % self.sample_interval == 0:
recon, mu, logvar = self.net(fixed_imgs)
recon = recon.view(-1, self.in_channel, self.image_size,
self.image_size)
x_hat_path = os.path.join(
self.sample_folder,
'output_epoch{}.png'.format(epoch + 1))
save_image(recon, x_hat_path, normalize=True)
print("[*] Save sample images!")
step += 1
if (epoch + 1) % self.ckpt_interval == 0:
ckpt_path = os.path.join(self.ckpt_folder,
'ckpt_epoch{}.pth'.format(epoch + 1))
torch.save(self.net.state_dict(), ckpt_path)
print("[*] Checkpoint saved!")
print("[*] Learning finished!")
ckpt_path = os.path.join(self.ckpt_folder, 'final_model.pth')
torch.save(self.net.state_dict(), ckpt_path)
print("[*] Final weight saved!")
help='output result for tensorboard and model checkpoint')
parser.add_argument('--cuda', action='store_true')
args = parser.parse_args()
# load checkpoint
if args.resume is not None:
checkpoint = torch.load(args.resume,
map_location=lambda storage, loc: storage)
print("checkpoint loaded!")
print("val loss: {}\tepoch: {}\t".format(checkpoint['val_loss'],
checkpoint['epoch']))
# model
model = VAE(args.image_size)
if args.resume is not None:
model.load_state_dict(checkpoint['state_dict'])
# criterion
criterion = VAELoss(size_average=True, kl_weight=args.kl_weight)
if args.cuda is True:
model = model.cuda()
criterion = criterion.cuda()
# load data
train_loader, val_loader = load_vae_train_datasets(input_size=args.image_size,
data=args.data,
batch_size=args.batch_size)
# load optimizer and scheduler
opt = torch.optim.Adam(params=model.parameters(),
lr=args.lr,
def test():
parser = argparse.ArgumentParser(description='Train VAE.')
parser.add_argument('-c', '--config', help='Config file.')
args = parser.parse_args()
print(args)
c = json.load(open(args.config))
print(c)
# clear param store
pyro.clear_param_store()
lookback = 50
test_start_date = datetime.strptime(
c['test_start_date'], '%Y/%m/%d') if c['test_start_date'] else None
test_end_date = datetime.strptime(
c['test_end_date'], '%Y/%m/%d') if c['test_end_date'] else None
min_sequence_length_test = 2 * (c['series_length'] + lookback)
max_n_files = None
out_path = Path(c['out_dir'])
out_path.mkdir(exist_ok=True)
dataset_test = create_ticker_dataset(
c['in_dir'],
c['series_length'],
lookback,
min_sequence_length_test,
start_date=test_start_date,
end_date=test_end_date,
fixed_start_date=True,
normalised_returns=c['normalised_returns'],
max_n_files=max_n_files)
test_loader = DataLoader(dataset_test,
batch_size=c['batch_size'],
shuffle=False,
num_workers=0,
drop_last=True)
N_test_data = len(dataset_test)
print(f'N_test_data: {N_test_data}')
# setup the VAE
vae = VAE(c['series_length'], z_dim=c['z_dim'], use_cuda=c['cuda'])
if c['checkpoint_load']:
checkpoint = torch.load(c['checkpoint_load'])
vae.load_state_dict(checkpoint['model_state_dict'])
# optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
if 1:
find_similar(vae, test_loader, c['cuda'])
# Visualise first batch.
batch = next(iter(test_loader))
x = batch['series']
if c['cuda']:
x = x.cuda()
x = x.float()
x_reconst = vae.reconstruct_img(x)
x = x.cpu().numpy()
x_reconst = x_reconst.cpu().detach().numpy()
n = min(5, x.shape[0])
fig, axes = plt.subplots(n, 1, squeeze=False)
for s in range(n):
ax = axes[s, 0]
ax.plot(x[s])
ax.plot(x_reconst[s])
fig.savefig(out_path / f'test_batch.png')
############### training #########################
lr = 0.01
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, net.parameters()), lr=lr, weight_decay = 0.0001)
def adjust_lr(optimizer, decay_rate=0.95):
for param_group in optimizer.param_groups:
param_group['lr'] *= decay_rate
retrain = True
if os.path.exists(save_name):
print("Model parameters have already been trained before. Retrain ? (y/n)")
ans = input()
if not (ans == 'y'):
checkpoint = torch.load(save_name, map_location = device)
net.load_state_dict(checkpoint["net"])
optimizer.load_state_dict(checkpoint["optimizer"])
for g in optimizer.param_groups:
g['lr'] = lr
max_epochs = 1000
early_stop = EarlyStop(patience = 20, save_name = save_name)
net = net.to(device)
print("training on ", device)
for epoch in range(max_epochs):
train_loss, n, start = 0.0, 0, time.time()
for X, _ in tqdm.tqdm(train_iter, ncols = 50):
X = X.to(device)
X_hat, mean, logvar = net(X)
