def test_error_on_nan(self):
"""Test if the correct error is raised for non-finite loss values."""
model = TransE(triples_factory=self.triples_factory)
patience = 2
class NaNTrainingLoop(self.cls):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.patience = patience
def _process_batch(self, *args, **kwargs):
loss = super()._process_batch(*args, **kwargs)
self.patience -= 1
if self.patience < 0:
return torch.as_tensor([float("nan")], device=loss.device, dtype=torch.float32)
return loss
training_loop = NaNTrainingLoop(
model=model,
triples_factory=self.triples_factory,
optimizer=self.optimizer_cls(model.get_grad_params()),
)
with self.assertRaises(NonFiniteLossError):
training_loop.train(
triples_factory=self.triples_factory,
num_epochs=patience + 1,
batch_size=self.batch_size,
)
def setUp(self) -> None:
"""Set up the test case with a triples factory and TransE as an example model."""
self.batch_size = 16
self.embedding_dim = 8
self.factory = Nations().training
self.model = TransE(self.factory,
embedding_dim=self.embedding_dim).to_device_()
def pre_setup_hook(self) -> None:
"""Prepare case-level variables before the setup() function."""
self.triples_factory = Nations().training
self.loss = self.loss_cls()
self.model = TransE(triples_factory=self.triples_factory,
loss=self.loss,
random_seed=self.random_seed)
self.optimizer = self.optimizer_cls(self.model.get_grad_params())
class TestBaseModel(unittest.TestCase):
"""Test models are set in the right mode at the right time."""
batch_size: int
embedding_dim: int
factory: TriplesFactory
model: Model
def setUp(self) -> None:
"""Set up the test case with a triples factory and TransE as an example model."""
self.batch_size = 16
self.embedding_dim = 8
self.factory = Nations().training
self.model = TransE(self.factory, embedding_dim=self.embedding_dim).to_device_()
def _check_scores(self, scores) -> None:
"""Check the scores produced by a forward function."""
# check for finite values by default
assert torch.all(torch.isfinite(scores)).item()
def test_predict_scores_all_heads(self) -> None:
"""Test ``BaseModule.predict_scores_all_heads``."""
batch = torch.zeros(self.batch_size, 2, dtype=torch.long, device=self.model.device)
# Set into training mode to check if it is correctly set to evaluation mode.
self.model.train()
scores = self.model.predict_h(batch)
assert scores.shape == (self.batch_size, self.model.num_entities)
self._check_scores(scores)
assert not self.model.training
def test_predict_scores_all_tails(self) -> None:
"""Test ``BaseModule.predict_scores_all_tails``."""
batch = torch.zeros(self.batch_size, 2, dtype=torch.long, device=self.model.device)
# Set into training mode to check if it is correctly set to evaluation mode.
self.model.train()
scores = self.model.predict_t(batch)
assert scores.shape == (self.batch_size, self.model.num_entities)
self._check_scores(scores)
assert not self.model.training
def test_predict_scores(self) -> None:
"""Test ``BaseModule.predict_scores``."""
batch = torch.zeros(self.batch_size, 3, dtype=torch.long, device=self.model.device)
# Set into training mode to check if it is correctly set to evaluation mode.
self.model.train()
scores = self.model.predict_hrt(batch)
assert scores.shape == (self.batch_size, 1)
self._check_scores(scores)
assert not self.model.training
def test_checkpoints(self):
"""Test whether interrupting the given training loop type can be resumed using checkpoints."""
# Train a model in one shot
model = TransE(
triples_factory=self.triples_factory,
random_seed=self.random_seed,
)
training_loop = self._with_model(model)
losses = training_loop.train(
triples_factory=self.triples_factory,
num_epochs=self.num_epochs,
batch_size=self.batch_size,
use_tqdm=False,
use_tqdm_batch=False,
)
with tempfile.TemporaryDirectory() as directory:
name = "checkpoint.pt"
# Train a model for the first half
model = TransE(
triples_factory=self.triples_factory,
random_seed=self.random_seed,
)
training_loop = self._with_model(model)
training_loop.train(
triples_factory=self.triples_factory,
num_epochs=int(self.num_epochs // 2),
batch_size=self.batch_size,
checkpoint_name=name,
checkpoint_directory=directory,
checkpoint_frequency=0,
)
# Continue training of the first part
model = TransE(
triples_factory=self.triples_factory,
random_seed=123,
)
training_loop = self._with_model(model)
losses_2 = training_loop.train(
triples_factory=self.triples_factory,
num_epochs=self.num_epochs,
batch_size=self.batch_size,
checkpoint_name=name,
checkpoint_directory=directory,
checkpoint_frequency=0,
)
self.assertEqual(losses, losses_2)
def setUp(self) -> None:
"""Set up the test case."""
# Settings
self.batch_size = 8
self.embedding_dim = 7
# Initialize evaluator
self.evaluator = self.evaluator_cls(**(self.evaluator_kwargs or {}))
# Use small test dataset
self.factory = Nations().training
# Use small model (untrained)
self.model = TransE(triples_factory=self.factory, embedding_dim=self.embedding_dim)
def test_early_stopping(self):
"""Tests early stopping."""
# Set automatic_memory_optimization to false during testing
nations = Nations()
model: Model = TransE(triples_factory=nations.training)
evaluator = RankBasedEvaluator(automatic_memory_optimization=False)
stopper = EarlyStopper(
model=model,
evaluator=evaluator,
training_triples_factory=nations.training,
evaluation_triples_factory=nations.validation,
patience=self.patience,
relative_delta=self.relative_delta,
metric='mean_rank',
)
training_loop = SLCWATrainingLoop(
model=model,
triples_factory=nations.training,
optimizer=Adam(params=model.get_grad_params()),
automatic_memory_optimization=False,
)
losses = training_loop.train(
triples_factory=nations.training,
num_epochs=self.max_num_epochs,
batch_size=self.batch_size,
stopper=stopper,
use_tqdm=False,
)
self.assertEqual(stopper.number_results,
len(losses) // stopper.frequency)
self.assertEqual(self.stop_epoch,
len(losses),
msg='Did not stop early like it should have')
def test_error_on_nan(self):
"""Test if the correct error is raised for non-finite loss values."""
model = TransE(triples_factory=self.triples_factory)
training_loop = NaNTrainingLoop(model=model, patience=2)
with self.assertRaises(NonFiniteLossError):
training_loop.train(num_epochs=3, batch_size=self.batch_size)
def _test_checkpoints(self, training_loop_type: str):
"""Test whether interrupting the given training loop type can be resumed using checkpoints."""
training_loop_class = get_training_loop_cls(training_loop_type)
# Train a model in one shot
model = TransE(
triples_factory=self.triples_factory,
random_seed=self.random_seed,
)
optimizer_cls = get_optimizer_cls(None)
optimizer = optimizer_cls(params=model.get_grad_params())
training_loop = training_loop_class(model=model, optimizer=optimizer, automatic_memory_optimization=False)
losses = training_loop.train(
num_epochs=self.num_epochs,
batch_size=self.batch_size,
use_tqdm=False,
use_tqdm_batch=False,
)
# Train a model for the first half
model = TransE(
triples_factory=self.triples_factory,
random_seed=self.random_seed,
)
optimizer_cls = get_optimizer_cls(None)
optimizer = optimizer_cls(params=model.get_grad_params())
training_loop = training_loop_class(model=model, optimizer=optimizer, automatic_memory_optimization=False)
training_loop.train(
num_epochs=int(self.num_epochs // 2),
batch_size=self.batch_size,
checkpoint_name=self.checkpoint_file,
checkpoint_directory=self.temporary_directory.name,
checkpoint_frequency=0,
)
# Continue training of the first part
model = TransE(
triples_factory=self.triples_factory,
random_seed=123,
)
optimizer_cls = get_optimizer_cls(None)
optimizer = optimizer_cls(params=model.get_grad_params())
training_loop = training_loop_class(model=model, optimizer=optimizer, automatic_memory_optimization=False)
losses_2 = training_loop.train(
num_epochs=self.num_epochs,
batch_size=self.batch_size,
checkpoint_name=self.checkpoint_file,
checkpoint_directory=self.temporary_directory.name,
checkpoint_frequency=0,
)
self.assertEqual(losses, losses_2)
def test_blacklist_loss_on_slcwa(self):
"""Test an allowed sLCWA loss."""
model = TransE(
triples_factory=self.triples_factory,
loss=CrossEntropyLoss(),
)
with self.assertRaises(TrainingApproachLossMismatchError):
NaNTrainingLoop(model=model, patience=2, automatic_memory_optimization=False)
def test_sub_batching(self):
"""Test if sub-batching works as expected."""
model = TransE(triples_factory=self.triples_factory)
training_loop = DummyTrainingLoop(
model=model,
sub_batch_size=self.sub_batch_size,
automatic_memory_optimization=False,
)
training_loop.train(num_epochs=1, batch_size=self.batch_size, sub_batch_size=self.sub_batch_size)
def test_lcwa_margin_ranking_loss_helper(self):
"""Test if output is correct for the LCWA training loop use case."""
factory = TriplesFactory.from_labeled_triples(triples=self.triples)
loss_cls = MarginRankingLoss(
margin=0,
reduction='sum',
)
model = TransE(
triples_factory=factory,
embedding_dim=8,
preferred_device='cpu',
loss=loss_cls,
)
loop = LCWATrainingLoop(model=model, triples_factory=factory)
loss = loop._mr_loss_helper(predictions=self.predictions,
labels=self.labels)
self.assertEqual(14, loss)
loss_cls = MarginRankingLoss(
margin=0,
reduction='mean',
)
model = TransE(
triples_factory=factory,
embedding_dim=8,
preferred_device='cpu',
loss=loss_cls,
)
loop = LCWATrainingLoop(model=model, triples_factory=factory)
loss = loop._mr_loss_helper(predictions=self.predictions,
labels=self.labels)
self.assertEqual(1, loss)
def _build_model(
self,
triples_factory: CoreTriplesFactory
) -> TransE:
"""Build new TransE model for embedding.
Parameters
------------------
graph: Graph
The graph to build the model for.
"""
return TransE(
triples_factory=triples_factory,
embedding_dim=self._embedding_size,
scoring_fct_norm=self._scoring_fct_norm
)
class TrainingLoopTestCase(unittest_templates.GenericTestCase[TrainingLoop]):
"""A generic test case for training loops."""
model: Model
factory: TriplesFactory
loss_cls: ClassVar[Type[Loss]]
loss: Loss
optimizer_cls: ClassVar[Type[Optimizer]] = Adam
optimizer: Optimizer
random_seed = 0
batch_size: int = 128
sub_batch_size: int = 30
num_epochs: int = 10
def pre_setup_hook(self) -> None:
"""Prepare case-level variables before the setup() function."""
self.triples_factory = Nations().training
self.loss = self.loss_cls()
self.model = TransE(triples_factory=self.triples_factory, loss=self.loss, random_seed=self.random_seed)
self.optimizer = self.optimizer_cls(self.model.get_grad_params())
def _with_model(self, model: Model) -> TrainingLoop:
return self.cls(
model=model,
triples_factory=self.triples_factory,
automatic_memory_optimization=False,
optimizer=self.optimizer_cls(model.get_grad_params()),
)
def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]: # noqa: D102
kwargs = super()._pre_instantiation_hook(kwargs=kwargs)
kwargs["triples_factory"] = self.triples_factory
kwargs["automatic_memory_optimization"] = False
kwargs["optimizer"] = self.optimizer
kwargs["model"] = self.model
return kwargs
def test_train(self):
"""Test training does not error."""
self.instance.train(
triples_factory=self.triples_factory,
num_epochs=1,
)
def test_sub_batching(self):
"""Test if sub-batching works as expected."""
self.instance.train(
triples_factory=self.triples_factory,
num_epochs=1,
batch_size=self.batch_size,
sub_batch_size=self.sub_batch_size,
)
def test_sub_batching_support(self):
"""Test if sub-batching works as expected."""
model = ConvE(triples_factory=self.triples_factory)
training_loop = self._with_model(model)
with self.assertRaises(NotImplementedError):
training_loop.train(
triples_factory=self.triples_factory,
num_epochs=1,
batch_size=self.batch_size,
sub_batch_size=self.sub_batch_size,
)
def test_error_on_nan(self):
"""Test if the correct error is raised for non-finite loss values."""
model = TransE(triples_factory=self.triples_factory)
patience = 2
class NaNTrainingLoop(self.cls):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.patience = patience
def _process_batch(self, *args, **kwargs):
loss = super()._process_batch(*args, **kwargs)
self.patience -= 1
if self.patience < 0:
return torch.as_tensor([float("nan")], device=loss.device, dtype=torch.float32)
return loss
training_loop = NaNTrainingLoop(
model=model,
triples_factory=self.triples_factory,
optimizer=self.optimizer_cls(model.get_grad_params()),
)
with self.assertRaises(NonFiniteLossError):
training_loop.train(
triples_factory=self.triples_factory,
num_epochs=patience + 1,
batch_size=self.batch_size,
)
def test_checkpoints(self):
"""Test whether interrupting the given training loop type can be resumed using checkpoints."""
# Train a model in one shot
model = TransE(
triples_factory=self.triples_factory,
random_seed=self.random_seed,
)
training_loop = self._with_model(model)
losses = training_loop.train(
triples_factory=self.triples_factory,
num_epochs=self.num_epochs,
batch_size=self.batch_size,
use_tqdm=False,
use_tqdm_batch=False,
)
with tempfile.TemporaryDirectory() as directory:
name = "checkpoint.pt"
# Train a model for the first half
model = TransE(
triples_factory=self.triples_factory,
random_seed=self.random_seed,
)
training_loop = self._with_model(model)
training_loop.train(
triples_factory=self.triples_factory,
num_epochs=int(self.num_epochs // 2),
batch_size=self.batch_size,
checkpoint_name=name,
checkpoint_directory=directory,
checkpoint_frequency=0,
)
# Continue training of the first part
model = TransE(
triples_factory=self.triples_factory,
random_seed=123,
)
training_loop = self._with_model(model)
losses_2 = training_loop.train(
triples_factory=self.triples_factory,
num_epochs=self.num_epochs,
batch_size=self.batch_size,
checkpoint_name=name,
checkpoint_directory=directory,
checkpoint_frequency=0,
)
self.assertEqual(losses, losses_2)
class _AbstractEvaluatorTests:
"""A test case for quickly defining common tests for evaluators models."""
# The triples factory and model
factory: TriplesFactory
model: Model
#: The evaluator to be tested
evaluator_cls: ClassVar[Type[Evaluator]]
evaluator_kwargs: ClassVar[Optional[Mapping[str, Any]]] = None
# Settings
batch_size: int
embedding_dim: int
#: The evaluator instantiation
evaluator: Evaluator
def setUp(self) -> None:
"""Set up the test case."""
# Settings
self.batch_size = 8
self.embedding_dim = 7
# Initialize evaluator
self.evaluator = self.evaluator_cls(**(self.evaluator_kwargs or {}))
# Use small test dataset
self.factory = Nations().training
# Use small model (untrained)
self.model = TransE(triples_factory=self.factory,
embedding_dim=self.embedding_dim)
def _get_input(
self,
inverse: bool = False,
) -> Tuple[torch.LongTensor, torch.FloatTensor,
Optional[torch.BoolTensor]]:
# Get batch
hrt_batch = self.factory.mapped_triples[:self.batch_size].to(
self.model.device)
# Compute scores
if inverse:
scores = self.model.score_h(rt_batch=hrt_batch[:, 1:])
else:
scores = self.model.score_t(hr_batch=hrt_batch[:, :2])
# Compute mask only if required
if self.evaluator.requires_positive_mask:
# TODO: Re-use filtering code
triples = self.factory.mapped_triples
if inverse:
sel_col, start_col = 0, 1
else:
sel_col, start_col = 2, 0
stop_col = start_col + 2
# shape: (batch_size, num_triples)
triple_mask = (triples[None, :, start_col:stop_col] ==
hrt_batch[:, None, start_col:stop_col]).all(dim=-1)
batch_indices, triple_indices = triple_mask.nonzero(as_tuple=True)
entity_indices = triples[triple_indices, sel_col]
# shape: (batch_size, num_entities)
mask = torch.zeros_like(scores, dtype=torch.bool)
mask[batch_indices, entity_indices] = True
else:
mask = None
return hrt_batch, scores, mask
def test_process_tail_scores_(self) -> None:
"""Test the evaluator's ``process_tail_scores_()`` function."""
hrt_batch, scores, mask = self._get_input()
true_scores = scores[torch.arange(0, hrt_batch.shape[0]),
hrt_batch[:, 2]][:, None]
self.evaluator.process_tail_scores_(
hrt_batch=hrt_batch,
true_scores=true_scores,
scores=scores,
dense_positive_mask=mask,
)
def test_process_head_scores_(self) -> None:
"""Test the evaluator's ``process_head_scores_()`` function."""
hrt_batch, scores, mask = self._get_input(inverse=True)
true_scores = scores[torch.arange(0, hrt_batch.shape[0]),
hrt_batch[:, 0]][:, None]
self.evaluator.process_head_scores_(
hrt_batch=hrt_batch,
true_scores=true_scores,
scores=scores,
dense_positive_mask=mask,
)
def test_finalize(self) -> None:
# Process one batch
hrt_batch, scores, mask = self._get_input()
true_scores = scores[torch.arange(0, hrt_batch.shape[0]),
hrt_batch[:, 2]][:, None]
self.evaluator.process_tail_scores_(
hrt_batch=hrt_batch,
true_scores=true_scores,
scores=scores,
dense_positive_mask=mask,
)
result = self.evaluator.finalize()
assert isinstance(result, MetricResults)
self._validate_result(result=result,
data={
'batch': hrt_batch,
'scores': scores,
'mask': mask
})
def _validate_result(
self,
result: MetricResults,
data: Dict[str, torch.Tensor],
):
logger.warning(
f'{self.__class__.__name__} did not overwrite _validate_result.')
from pykeen.triples import TriplesFactory
from pykeen.evaluation import RankBasedEvaluator
from pykeen.models import TransE
from torch.optim import Adam
from pykeen.training import SLCWATrainingLoop
n_tokeep = 300
minimum = 500
tf = TriplesFactory.from_path(f'data/rare/rare_{minimum}_{n_tokeep}.csv')
training, testing, validation = tf.split([.8, .1, .1])
# Pick model, optimizer, training approach
model = TransE(triples_factory=training)
optimizer = Adam(params=model.get_grad_params())
training_loop = SLCWATrainingLoop(
model=model,
triples_factory=training,
optimizer=optimizer,
)
# train
training_loop.train(
triples_factory=training,
num_epochs=500,
batch_size=256,
)
# evaluate
evaluator = RankBasedEvaluator(ks=[50])
mapped_triples = testing.mapped_triples