def transfer_batch_to_device(self, batch: PathContextBatch,
device: torch.device) -> PathContextBatch:
# Dict str -> torch.Tensor [seq length; batch size * n_context]
for k in batch.context:
batch.context[k] = batch.context[k].to(self.device)
# [seq length; batch size]
batch.labels = batch.labels.to(self.device)
return batch
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 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)
def transfer_batch_to_device(self, batch: PathContextBatch,
device: torch.device) -> PathContextBatch:
batch.move_to_device(device)
return batch
def collate_wrapper(batch: List[PathContextSample]) -> PathContextBatch:
return PathContextBatch(batch)