dicee.trainer.model_parallelism

Tensor Parallelism trainer for ensemble knowledge graph embeddings.

This module implements the tensor parallelism training strategy described in:

“Multiple Run Ensemble Learning with Low-Dimensional Knowledge Graph Embeddings”: https://arxiv.org/abs/2104.05003

The TensorParallel trainer creates an ensemble of models, each trained on a separate GPU, allowing efficient utilization of multi-GPU systems through model parallelism.

Attributes

logger

Classes

TensorParallel

Abstract base class for KGE model trainers.

Functions

`extract_input_outputs`(z[, device])
`forward_backward_update_loss`(→ float)

Module Contents

dicee.trainer.model_parallelism.logger

dicee.trainer.model_parallelism.extract_input_outputs(z: list, device=None)

dicee.trainer.model_parallelism.forward_backward_update_loss(z: Tuple, ensemble_model) → float

class dicee.trainer.model_parallelism.TensorParallel(args, callbacks)

Bases: dicee.abstracts.AbstractTrainer

Abstract base class for KGE model trainers.

Provides the callback dispatch mechanism shared by all concrete trainer implementations (TorchTrainer, TorchDDPTrainer, etc.). Sub-classes call the on_* hooks at the appropriate points in the training loop so that any registered AbstractCallback can react.

Parameters:

args (argparse.Namespace or similar) – Processed configuration object. Must expose at least random_seed (int).
callbacks (list of AbstractCallback) – Ordered list of callback instances to invoke at each lifecycle hook.

fit(*args, **kwargs): Train model