model_maintainer#

Module Contents#

ModelMaintainer

Maintain PyTorch model.

SerialModelMaintainer

"Maintain PyTorch model.
class ModelMaintainer(model: torch.nn.Module, cuda: bool, device: str = None)#

Bases: object

Maintain PyTorch model.

Provide necessary attributes and operation methods. More features with local or global model will be implemented here.

Parameters:

  • model (torch.nn.Module) – PyTorch model.

  • cuda (bool) – Use GPUs or not.

  • device (str, optional) – Assign model/data to the given GPUs. E.g., ‘device:0’ or ‘device:0,1’. Defaults to None. If device is None and cuda is True, FedLab will set the gpu with the largest memory as default.

property model: torch.nn.Module#

Return torch.nn.module.

property model_parameters: torch.Tensor#

Return serialized model parameters.

property model_gradients: torch.Tensor#

Return serialized model gradients.

property shape_list: List[torch.Tensor]#

Return shape of model parameters.

Currently, this attributes used in tensor compression.

set_model(parameters: torch.Tensor)#

Assign parameters to self._model.

class SerialModelMaintainer(model: torch.nn.Module, num_clients: int, cuda: bool, device: str = None, personal: bool = False)#

Bases: ModelMaintainer

“Maintain PyTorch model.

Provide necessary attributes and operation methods. More features with local or global model will be implemented here.

Parameters:

  • model (torch.nn.Module) – PyTorch model.

  • num_clients (int) – The number of independent models.

  • cuda (bool) – Use GPUs or not.

  • device (str, optional) – Assign model/data to the given GPUs. E.g., ‘device:0’ or ‘device:0,1’. Defaults to None. If device is None and cuda is True, FedLab will set the gpu with the largest idle memory as default.

  • personal (bool, optional) – If Ture is passed, SerialModelMaintainer will generate the copy of local parameters list and maintain them respectively. These paremeters are indexed by [0, num-1]. Defaults to False.

set_model(parameters: torch.Tensor = None, id: int = None)#

Assign parameters to self._model.

Note

parameters and id can not be None at the same time. If id is None, this function load the given parameters. If id is not None, this function load the parameters of given id first and the parameters attribute will be ignored.

Parameters:

  • parameters (torch.Tensor, optional) – Model parameters. Defaults to None.

  • id (int, optional) – Load the model parameters of client id. Defaults to None.