class CombinedRegularizerTest(_RegularizerTestCase, unittest.TestCase):
"""Test the combined regularizer."""
regularizer_cls = CombinedRegularizer
regularizer_kwargs = {
'regularizers': [
LpRegularizer(weight=0.1, p=1, device=resolve_device()),
LpRegularizer(weight=0.7, p=2, device=resolve_device()),
],
}
def _expected_penalty(self, x: torch.FloatTensor) -> torch.FloatTensor: # noqa: D102
regularizers = self.regularizer_kwargs['regularizers']
return sum(r.weight * r.forward(x) for r in regularizers) / sum(r.weight for r in regularizers)
def setUpClass(cls):
"""Set up a shared result."""
cls.device = resolve_device('cuda')
cls.dataset = Nations()
cls.model = MockModel(triples_factory=cls.dataset.training)
# The MockModel gives the highest score to the highest entity id
max_score = cls.dataset.num_entities - 1
# The test triples are created to yield the third highest score on both head and tail prediction
cls.dataset.testing.mapped_triples = torch.tensor(
[[max_score - 2, 0, max_score - 2]])
# Write new mapped triples to the model, since the model's triples will be used to filter
# These triples are created to yield the highest score on both head and tail prediction for the
# test triple at hand
cls.dataset.training.mapped_triples = torch.tensor([
[max_score - 2, 0, max_score],
[max_score, 0, max_score - 2],
], )
# The validation triples are created to yield the second highest score on both head and tail prediction for the
# test triple at hand
cls.dataset.validation.mapped_triples = torch.tensor([
[max_score - 2, 0, max_score - 1],
[max_score - 1, 0, max_score - 2],
], )
def setUp(self):
"""Prepare for testing the scoring functions."""
self.generator = torch.random.manual_seed(seed=42)
self.triples_factory = MinimalTriplesFactory
self.device = resolve_device()
self.model = SimpleInteractionModel(
triples_factory=self.triples_factory).to(self.device)
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.device = resolve_device("cpu")
self.model = TransE(triples_factory=self.factory,
embedding_dim=self.embedding_dim).to(self.device)
def setUp(self) -> None:
"""Set up the test case with a triples factory and model."""
self.device = resolve_device()
self.triples_factory = Nations().training
self.batch_size = 16
self.positive_batch = self.triples_factory.mapped_triples[:self.batch_size, :].to(device=self.device)
super().setUp()
# move test instance to device
self.instance = self.instance.to(self.device)
def setUp(self) -> None:
"""Set up the test case with a triples factory and model."""
self.generator = torch.random.manual_seed(seed=42)
self.batch_size = 16
self.triples_factory = Nations().training
self.device = resolve_device()
self.regularizer = self.regularizer_cls(
device=self.device,
**(self.regularizer_kwargs or {}),
)
self.positive_batch = self.triples_factory.mapped_triples[:self.batch_size, :].to(device=self.device)
def setUp(self) -> None:
"""Set up the test case."""
self.generator = torch.random.manual_seed(seed=42)
self.device = resolve_device()
self.kwargs = {"weight": 0.5, "epsilon": 1e-5}
self.instance = self.cls(
**(self.kwargs or {}),
).to(self.device)
self.num_entities = 10
self.num_relations = 5
self.entities_weight = rand(self.num_entities, 10, generator=self.generator, device=self.device)
self.relations_weight = rand(self.num_relations, 20, generator=self.generator, device=self.device)
self.normal_vector_weight = rand(self.num_relations, 20, generator=self.generator, device=self.device)
def setUp(self) -> None:
"""Set up the test case."""
self.generator = torch.random.manual_seed(seed=42)
self.device = resolve_device()
self.regularizer_kwargs = {'weight': .5, 'epsilon': 1e-5}
self.regularizer = TransHRegularizer(
device=self.device,
**(self.regularizer_kwargs or {}),
)
self.num_entities = 10
self.num_relations = 5
self.entities_weight = torch.rand(self.num_entities, 10, device=self.device, generator=self.generator)
self.relations_weight = torch.rand(self.num_relations, 20, device=self.device, generator=self.generator)
self.normal_vector_weight = torch.rand(self.num_relations, 20, device=self.device, generator=self.generator)
def setUpClass(cls):
"""Set up a shared result."""
cls.device = resolve_device('cuda')
cls.result = pipeline(
model='TransE',
dataset='nations',
training_kwargs=dict(num_epochs=5, use_tqdm=False),
evaluation_kwargs=dict(use_tqdm=False),
device=cls.device,
random_seed=42,
)
cls.model = cls.result.model
nations = Nations()
cls.testing_mapped_triples = nations.testing.mapped_triples.to(cls.model.device)
def get_model_size( # noqa: C901
*,
dataset: Union[None, str, Type[Dataset]] = None,
dataset_kwargs: Optional[Mapping[str, Any]] = None,
training: Optional[TriplesFactory] = None,
testing: Optional[TriplesFactory] = None,
validation: Optional[TriplesFactory] = None,
model: Union[str, Type[Model]],
model_kwargs: Optional[Mapping[str, Any]] = None,
loss: Union[None, str, Type[Loss]] = None,
loss_kwargs: Optional[Mapping[str, Any]] = None,
regularizer: Union[None, str, Type[Regularizer]] = None,
regularizer_kwargs: Optional[Mapping[str, Any]] = None,
**_kwargs,
) -> int:
"""Make a model instance, similarly to how the pipelin is started, then return the model size."""
device = resolve_device('cpu')
dataset = get_dataset(
dataset=dataset,
dataset_kwargs=dataset_kwargs,
training=training,
testing=testing,
validation=validation,
)
if model_kwargs is None:
model_kwargs = {}
if regularizer is not None:
regularizer_cls = get_regularizer_cls(regularizer)
model_kwargs['regularizer'] = regularizer_cls(
device=device,
**(regularizer_kwargs or {}),
)
if loss is not None:
loss_cls = get_loss_cls(loss)
model_kwargs['loss'] = loss_cls(**(loss_kwargs or {}))
model = get_model_cls(model)
model_instance: Model = model(
random_seed=0,
preferred_device=device,
triples_factory=dataset.training,
**model_kwargs,
)
return model_instance.num_parameter_bytes
def setUp(self) -> None:
"""Set up the test case."""
self.generator = torch.random.manual_seed(seed=42)
self.device = resolve_device()
