def train(): parser = argparse.ArgumentParser(description='Train VAE.') parser.add_argument('-c', '--config', default='train_config.json', help='Config file.') args = parser.parse_args() print(args) c = json.load(open(args.config)) print(c) pyro.clear_param_store() # TODO: Move to config file. lookback = 50 max_n_files = None train_start_date = datetime.strptime(c['train_start_date'], '%Y/%m/%d') train_end_date = datetime.strptime(c['train_end_date'], '%Y/%m/%d') val_start_date = datetime.strptime(c['val_start_date'], '%Y/%m/%d') val_end_date = datetime.strptime(c['val_end_date'], '%Y/%m/%d') min_sequence_length_train = 2 * (c['series_length'] + lookback) min_sequence_length_test = 2 * (c['series_length'] + lookback) out_path = Path(c['out_dir']) out_path.mkdir(exist_ok=True) dataset_train = create_ticker_dataset(c['in_dir'], c['series_length'], lookback, min_sequence_length_train, start_date=train_start_date, end_date=train_end_date, normalised_returns=c['normalised_returns'], max_n_files=max_n_files) dataset_val = create_ticker_dataset(c['in_dir'], c['series_length'], lookback, min_sequence_length_test, start_date=val_start_date, end_date=val_end_date, fixed_start_date=True, normalised_returns=c['normalised_returns'], max_n_files=max_n_files) train_loader = DataLoader(dataset_train, batch_size=c['batch_size'], shuffle=True, num_workers=0, drop_last=True) val_loader = DataLoader(dataset_val, batch_size=c['batch_size'], shuffle=False, num_workers=0, drop_last=True) N_train_data = len(dataset_train) N_val_data = len(dataset_val) N_mini_batches = N_train_data // c['batch_size'] N_train_time_slices = c['batch_size'] * N_mini_batches print(f'N_train_data: {N_train_data}, N_val_data: {N_val_data}') # setup the VAE vae = VAE(c['series_length'], z_dim=c['z_dim'], hidden_dims=c['hidden_dims'], use_cuda=c['cuda']) # setup the optimizer adam_args = {"lr": c['learning_rate']} optimizer = Adam(adam_args) # setup the inference algorithm elbo = JitTrace_ELBO() if c['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']) train_elbo = [] val_elbo = [] # training loop for epoch in range(c['n_epochs']): # initialize loss accumulator epoch_loss = 0. # do a training epoch over each mini-batch x returned # by the data loader for batch in train_loader: x = batch['series'] # if on GPU put mini-batch into CUDA memory if c['cuda']: x = x.cuda() # do ELBO gradient and accumulate loss epoch_loss += svi.step(x.float()) # report training diagnostics normalizer_train = len(train_loader.dataset) total_epoch_loss_train = epoch_loss / normalizer_train train_elbo.append(total_epoch_loss_train) print("[epoch %03d] average training loss: %.4f" % (epoch, total_epoch_loss_train)) torch.save({ 'epoch': epoch, 'model_state_dict': vae.state_dict(), # 'optimizer_state_dict': optimizer.state_dict(), 'train_loss': epoch_loss }, out_path / c['checkpoint_save'].format(epoch)) if epoch % c['val_frequency'] == 0: # initialize loss accumulator val_loss = 0. # compute the loss over the entire test set for i, batch in enumerate(val_loader): x = batch['series'] # if on GPU put mini-batch into CUDA memory if c['cuda']: x = x.cuda() x = x.float() # compute ELBO estimate and accumulate loss val_loss += svi.evaluate_loss(x) if i == 0: # Visualise first batch. 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'val_{epoch:03d}.png') plt.close(fig) # report test diagnostics normalizer_val = len(val_loader.dataset) total_epoch_loss_val = val_loss / normalizer_val val_elbo.append(total_epoch_loss_val) print("[epoch %03d] average val loss: %.4f" % (epoch, total_epoch_loss_val)) # t-SNE. all_z_latents = [] for batch in val_loader: x = batch['series'] # z_latents = minibatch_inference(dmm, test_batch) # z_latents = encode_x_to_z(dmm, test_batch, sample_z_t=False) # x, z, x_reconst = test_minibatch(dmm, test_batch, args, sample_z=True) if c['cuda']: x = x.cuda() z_loc, z_scale, z = vae.encode_x(x.float()) all_z_latents.append(z.cpu().numpy()) # all_latents = torch.cat(all_z_latents, dim=0) all_latents = np.concatenate(all_z_latents, axis=0) # Run t-SNE with 2 output dimensions. from sklearn.manifold import TSNE model_tsne = TSNE(n_components=2, random_state=0) # z_states = all_latents.detach().cpu().numpy() z_states = all_latents z_embed = model_tsne.fit_transform(z_states) # Plot t-SNE embedding. fig = plt.figure() plt.scatter(z_embed[:, 0], z_embed[:, 1], s=10) fig.savefig(out_path / f'tsne_{epoch:03d}.png') plt.close(fig) print('Finished training.')
def train(): parser = argparse.ArgumentParser(description='Train VAE.') parser.add_argument('-c', '--config', default='train_config.json', help='Config file.') args = parser.parse_args() print(args) c = json.load(open(args.config)) print(c) # clear param store pyro.clear_param_store() input_dim = 1 max_n_examples = None out_path = Path(c['out_dir']) out_path.mkdir(exist_ok=True) if 0: dataset_train = create_ranking_dataset(c['training_filename'], 0.0, 0.7, max_n_examples=max_n_examples) dataset_val = create_ranking_dataset(c['training_filename'], 0.7, 1.0, max_n_examples=max_n_examples) train_loader = DataLoader(dataset_train, batch_size=c['batch_size'], shuffle=True, num_workers=3, drop_last=True) val_loader = DataLoader(dataset_val, batch_size=c['batch_size'], shuffle=False, num_workers=3, drop_last=True) else: dataset_train = create_ranking_dataset(n_examples=100) dataset_val = create_ranking_dataset(n_examples=100) train_loader = DataLoader(dataset_train, batch_size=8, shuffle=True, num_workers=3, drop_last=True) val_loader = DataLoader(dataset_val, batch_size=8, shuffle=False, num_workers=3, drop_last=True) c['num_features'] = 2 N_train_data = len(dataset_train) N_val_data = len(dataset_val) N_mini_batches = N_train_data // c['batch_size'] N_train_time_slices = c['batch_size'] * N_mini_batches print(f'N_train_data: {N_train_data}, N_val_data: {N_val_data}') # setup the VAE vae = VAE(c['num_features'], z_dim=c['z_dim'], use_cuda=c['cuda']) # setup the optimizer adam_args = {"lr": c['learning_rate']} optimizer = Adam(adam_args) # setup the inference algorithm elbo = JitTrace_ELBO() if c['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']) train_elbo = [] val_elbo = [] # training loop num_epochs = c['n_epochs'] for epoch in range(num_epochs): print(f'Starting epoch {epoch} of {num_epochs}.') # initialize loss accumulator epoch_loss = 0. # do a training epoch over each mini-batch x returned # by the data loader for batch_num, batch in enumerate(train_loader): print(f'Batch {batch_num} of {N_mini_batches}.') features_1 = batch['features_1'] features_2 = batch['features_2'] target_class = batch['target_class'] if c['cuda']: features_1 = features_1.cuda() features_2 = features_2.cuda() target_class = target_class.cuda() # do ELBO gradient and accumulate loss epoch_loss += svi.step(features_1.float(), features_2.float(), target_class.float()) # report training diagnostics normalizer_train = len(train_loader.dataset) total_epoch_loss_train = epoch_loss / normalizer_train train_elbo.append(total_epoch_loss_train) print("[epoch %03d] average training loss: %.4f" % (epoch, total_epoch_loss_train)) torch.save( { 'epoch': epoch, 'model_state_dict': vae.state_dict(), # 'optimizer_state_dict': optimizer.state_dict(), 'train_loss': epoch_loss }, out_path / c['checkpoint_save'].format(epoch)) if epoch % c['val_frequency'] == 0: print('Evaluation validation data.') val_loss = 0. true_positives = 0 num_val_examples = 0 # Compute the loss over the validation set. for i, batch in enumerate(val_loader): features_1 = batch['features_1'] features_2 = batch['features_2'] target_class = batch['target_class'] # if on GPU put mini-batch into CUDA memory if c['cuda']: features_1 = features_1.cuda() features_2 = features_2.cuda() target_class = target_class.cuda() features_1 = features_1.float() features_2 = features_2.float() target_class = target_class.float() # compute ELBO estimate and accumulate loss val_loss += svi.evaluate_loss(features_1, features_2, target_class) if 1: z_1_loc, z_1_scale, z_1_sample = vae.encode_x(features_1) z_2_loc, z_2_scale, z_2_sample = vae.encode_x(features_2) pred = torch.sigmoid(z_1_loc - z_2_loc) pred[pred >= 0.5] = 1 pred[pred < 0.5] = 0 target_class = torch.unsqueeze(target_class, dim=1) pred_target = torch.cat([pred, target_class], dim=1) # print(f'pred vs target: {pred_target}') pred = pred.cpu().detach().numpy() target_class = target_class.cpu().detach().numpy() pred_correct = (pred == target_class).astype(int) # print(pred_correct) # print(target_class.shape) pred_correct.flatten() true_positives += np.sum(pred_correct) num_val_examples += pred_correct.shape[0] # accuracy = np.sum(pred_correct) / pred_correct.shape[0] # print('accuracy:', accuracy) # return if 0: # Make some rank predictions. z_1_loc, z_1_scale, z_1_sample = vae.encode_x(x_1) z_2_loc, z_2_scale, z_2_sample = vae.encode_x(x_2) z_loc = torch.cat([z_1_loc, z_2_loc], dim=1) print(f'y: {y}, z_loc: {z_loc}') return if 0: if i == 0: # Visualise first batch. 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'val_{epoch:03d}.png') plt.close(fig) # report test diagnostics normalizer_val = len(val_loader.dataset) total_epoch_loss_val = val_loss / normalizer_val val_elbo.append(total_epoch_loss_val) print("[epoch %03d] average val loss: %.4f" % (epoch, total_epoch_loss_val)) accuracy = true_positives / num_val_examples print(f'Accuracy: {accuracy}') print('Finished training.')