def sample(model, _config, _run): model_eval_config = _config['generation'] cuda = _config['gpu']['cuda'] training_config = _config['training'] # Setup Generation results_dir = exp_dir_path('sample_trajectories') n_samples = model_eval_config['n_trajectory_samples'] # Create Dataset dataset = MultimodalDataset( data_dir=exp_dir_path('../../../scenarios/multimodal_dataset/data'), modalities=['trajectory'], batch_size=training_config['batch_size'], eval_samples=training_config['eval_samples'], validation_size=training_config['validation_size'], seed=training_config['seed']) trj_norm = dataset.get_traj_normalization() # Sample cross-modality from symbol for i in tqdm(range(0, 10)): sample_t = generate_class(model, i, n_samples, cuda) # Trajectory utils.save_trajectory(data=sample_t, path=results_dir, filename='image_class_' + str(i) + '.png', traj_norm=trj_norm, tmp_path=results_dir) ex.add_artifact(os.path.join(results_dir, 'image_class_' + str(i) + '.png'), name='image_class_' + str(i) + '.png')
def train(_config, _run): # Read configs training_config = _config['training'] device = torch.device("cuda" if _config['gpu']['cuda'] else "cpu") results_dir = log_dir_path('results') artifact_storage_interval = _config['model_debug'][ 'artifact_storage_interval'] # Set seeds torch.backends.cudnn.deterministic = True torch.manual_seed(training_config['seed']) np.random.seed(training_config['seed']) random.seed(training_config['seed']) torch.cuda.manual_seed(training_config['seed']) # Create Classifier model = Trajectory_AE(b_dim=training_config['b_dim']).to(device) epochs = training_config['epochs'] # Create Dataset dataset = MultimodalDataset( modalities=['trajectory'], data_dir=exp_dir_path('../../../scenarios/multimodal_dataset/data'), batch_size=training_config['batch_size'], eval_samples=10, validation_size=0.1, seed=training_config['seed']) trj_norm = dataset.get_traj_normalization() train_loader, val_loader = dataset.train_loader, dataset.val_loader test_loader = dataset.get_test_loader(bsize=20) # Training objects optimizer = torch.optim.Adam(model.parameters(), lr=training_config['learning_rate']) best_loss = sys.maxsize for epoch in range(1, epochs + 1): train_loss = train_epoch(model, train_loader, optimizer, epoch, cuda=_config['gpu']['cuda']) val_loss = test_epoch(model, val_loader, cuda=_config['gpu']['cuda']) _run.log_scalar('train_loss', train_loss) _run.log_scalar('val_loss', val_loss) # Best Loss is_best = val_loss < best_loss best_loss = min(val_loss, best_loss) record_checkpoint(model=model, loss=val_loss, best_loss=best_loss, optimizer=optimizer, epoch=epoch, is_best=is_best) if epoch % artifact_storage_interval == 0: # Data with torch.no_grad(): model.eval() data = next(iter(test_loader))[1].to(device) # Generate modalities out = model(data) save_trajectory(data=data, path=results_dir, filename='trj_og_e' + str(epoch) + '.png', traj_norm=trj_norm, tmp_path=results_dir) ex.add_artifact(os.path.join(results_dir, "trj_og_e" + str(epoch) + '.png'), name="trajectory_og_e" + str(epoch) + '.png') save_trajectory(data=out, path=results_dir, filename='trj_mod_e' + str(epoch) + '.png', traj_norm=trj_norm, tmp_path=results_dir) ex.add_artifact(os.path.join(results_dir, "trj_mod_e" + str(epoch) + '.png'), name="trajectory_recon_e" + str(epoch) + '.png') # Final Saving ex.add_artifact( os.path.join( log_dir_path('trained_models'), 'trajectory_ae_checkpoint.pth.tar'), name='trajectory_ae_last_checkpoint.pth.tar') ex.add_artifact( os.path.join( log_dir_path('trained_models'), 'best_trajectory_ae_model.pth.tar'), name='best_trajectory_ae_model.pth.tar')