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')
