def test_forward(self):
with initialize_config_dir(config_dir=get_config_directory()):
data_folder, dataset_name = get_test_data_info()
config = compose("main",
overrides=[
f"data_folder={data_folder}",
f"dataset.name={dataset_name}"
])
dataset_folder = join(config.data_folder, config.dataset.name)
vocabulary = Vocabulary.load_vocabulary(
join(dataset_folder, config.vocabulary_name))
data_file_path = join(
dataset_folder,
f"{config.dataset.name}.{config.train_holdout}.c2s")
dataset = PathContextDataset(data_file_path, config, vocabulary, False)
batch = PathContextBatch(
[dataset[i] for i in range(config.hyper_parameters.batch_size)])
model = PathEncoder(
config.encoder,
config.decoder.decoder_size,
len(vocabulary.token_to_id),
vocabulary.token_to_id[PAD],
len(vocabulary.node_to_id),
vocabulary.node_to_id[PAD],
)
output = model(batch.contexts)
true_shape = (sum(batch.contexts_per_label),
config.decoder.decoder_size)
self.assertTupleEqual(true_shape, output.shape)
def train(config: DictConfig):
filter_warnings()
print_config(config)
seed_everything(config.seed)
known_models = {"code2seq": get_code2seq, "code2class": get_code2class, "typed-code2seq": get_typed_code2seq}
if config.name not in known_models:
print(f"Unknown model: {config.name}, try on of {known_models.keys()}")
vocabulary = Vocabulary.load_vocabulary(join(config.data_folder, config.dataset.name, config.vocabulary_name))
model, data_module = known_models[config.name](config, vocabulary)
# define logger
wandb_logger = WandbLogger(
project=f"{config.name}-{config.dataset.name}", log_model=True, offline=config.log_offline
)
wandb_logger.watch(model)
# define model checkpoint callback
checkpoint_callback = ModelCheckpoint(
dirpath=wandb_logger.experiment.dir,
filename="{epoch:02d}-{val_loss:.4f}",
period=config.save_every_epoch,
save_top_k=-1,
)
upload_checkpoint_callback = UploadCheckpointCallback(wandb_logger.experiment.dir)
# define early stopping callback
early_stopping_callback = EarlyStopping(
patience=config.hyper_parameters.patience, monitor="val_loss", verbose=True, mode="min"
)
# define callback for printing intermediate result
print_epoch_result_callback = PrintEpochResultCallback("train", "val")
# use gpu if it exists
gpu = 1 if torch.cuda.is_available() else None
# define learning rate logger
lr_logger = LearningRateMonitor("step")
trainer = Trainer(
max_epochs=config.hyper_parameters.n_epochs,
gradient_clip_val=config.hyper_parameters.clip_norm,
deterministic=True,
check_val_every_n_epoch=config.val_every_epoch,
log_every_n_steps=config.log_every_epoch,
logger=wandb_logger,
gpus=gpu,
progress_bar_refresh_rate=config.progress_bar_refresh_rate,
callbacks=[
lr_logger,
early_stopping_callback,
checkpoint_callback,
upload_checkpoint_callback,
print_epoch_result_callback,
],
resume_from_checkpoint=config.resume_from_checkpoint,
)
trainer.fit(model=model, datamodule=data_module)
trainer.test()
def test_forward(self):
with initialize_config_dir(config_dir=get_config_directory()):
data_folder, dataset_name = get_test_data_info()
config = compose("main", overrides=[f"data_folder={data_folder}", f"dataset.name={dataset_name}"])
dataset_folder = join(config.data_folder, config.dataset.name)
vocabulary = Vocabulary.load_vocabulary(join(dataset_folder, config.vocabulary_name))
data_file_path = join(dataset_folder, f"{config.dataset.name}.{config.train_holdout}.c2s")
dataset = PathContextDataset(data_file_path, config, vocabulary, False)
batch = PathContextBatch([dataset[i] for i in range(config.hyper_parameters.batch_size)])
number_of_paths = sum(batch.contexts_per_label)
model = PathDecoder(config.decoder, len(vocabulary.label_to_id), 0, 0)
fake_encoder_output = torch.rand(number_of_paths, config.decoder.decoder_size)
output = model(fake_encoder_output, batch.contexts_per_label, config.dataset.target.max_parts)
true_shape = (
config.dataset.target.max_parts,
config.hyper_parameters.batch_size,
len(vocabulary.label_to_id),
)
self.assertTupleEqual(true_shape, output.shape)
