openfl.federated.task.task_runner.CoreTaskRunner

class openfl.federated.task.task_runner.CoreTaskRunner(**kwargs)

Bases: object

Federated Learning Task Runner Class.

Class Attributes:
  • kwargs (dict) – Additional parameters passed to the function.

  • TASK_REGISTRY (dict) – Registry of tasks.

  • training_round_completed (bool) – Flag indicating if a training round has been completed.

  • tensor_dict_split_fn_kwargs (dict) – Key word arguments for determining which parameters to hold out from aggregation.

  • required_tensorkeys_for_function (dict) – Required tensorkeys for all public functions in CoreTaskRunner.

  • logger (logging.Logger) – Logger object for logging events.

  • opt_treatment (str) – Treatment of current instance optimizer.

Methods

adapt_tasks

Prepare tasks for the collaborator.

get_required_tensorkeys_for_function

Get the required tensors for specified function that could be called as part of a task.

get_tensor_dict

Return the tensor dictionary.

get_train_data_size

Get the number of training examples.

get_valid_data_size

Get the number of examples.

initialize_tensorkeys_for_functions

Set the required tensors for all publicly accessible task methods.

rebuild_model

Parse tensor names and update weights of model.

reset_opt_vars

Reset optimizer variables.

set_data_loader

Register a data loader initialized with local data path.

set_framework_adapter

Set framework adapter.

set_logger

Set up the log object.

set_model_provider

Retrieve a model and an optimizer from the interface object.

set_optimizer_treatment

Change the treatment of current instance optimizer.

set_task_provider

Set task registry.

set_tensor_dict

Set the tensor dictionary.

adapt_tasks()

Prepare tasks for the collaborator.

Using functions from a task provider (deserialized interface object) and registered task contracts prepares callable tasks to be invoked by the collaborator.

Preparing includes conditional model rebuilding and filling output dicts with tensors for aggregation and storing in local DB.

There is an assumption that any training task accepts optimizer as one of the arguments, thus the model should be aggregated after such tasks.

Returns:

None

get_required_tensorkeys_for_function(func_name, **kwargs)

Get the required tensors for specified function that could be called as part of a task.

By default, this is just all of the layers and optimizer of the model.

Parameters:
  • func_name (str) – The function name.

  • **kwargs – Additional parameters to pass to the function.

Returns:

list – List of required TensorKey.

get_tensor_dict(with_opt_vars=False)

Return the tensor dictionary.

Parameters:

with_opt_vars (bool) – Return the tensor dictionary including the optimizer tensors (Default=False).

Returns:

dict – Tensor dictionary {**dict, **optimizer_dict}.

get_train_data_size()

Get the number of training examples.

It will be used for weighted averaging in aggregation.

Returns:

int – The number of training examples.

get_valid_data_size()

Get the number of examples.

It will be used for weighted averaging in aggregation.

Returns:

int – The number of validation examples.

initialize_tensorkeys_for_functions(with_opt_vars=False)

Set the required tensors for all publicly accessible task methods.

By default, this is just all of the layers and optimizer of the model. Custom tensors should be added to this function.

Parameters:

with_opt_vars (bool) – Specify if we also want to set the variables of the optimizer. Default is False.

Returns:

None

rebuild_model(input_tensor_dict, validation=False, device='cpu')

Parse tensor names and update weights of model. Handles the optimizer treatment.

Parameters:
  • input_tensor_dict (dict) – The input tensor dictionary.

  • validation (bool) – If True, perform validation. Default is False.

  • device (str) – The device to use. Default is ‘cpu’.

Returns:

None

reset_opt_vars()

Reset optimizer variables.

Returns:

None

set_data_loader(data_loader)

Register a data loader initialized with local data path.

Parameters:

data_loader – Data loader object.

Returns:

None

set_framework_adapter(framework_adapter)

Set framework adapter.

Setting a framework adapter allows first extraction of the weigths of the model with the purpose to make a list of parameters to be aggregated.

Parameters:

framework_adapter – Framework adapter object.

Returns:

None

set_logger()

Set up the log object.

Returns:

None

set_model_provider(model_provider)

Retrieve a model and an optimizer from the interface object.

Parameters:

model_provider – Model provider object.

Returns:

None

set_optimizer_treatment(opt_treatment)

Change the treatment of current instance optimizer.

Parameters:

opt_treatment (str) – The optimizer treatment.

Returns:

None

set_task_provider(task_provider)

Set task registry.

This method recieves Task Interface object as an argument and uses provided callables and information to prepare tasks that may be called by the collaborator component.

Parameters:

task_provider – Task provider object.

Returns:

None

set_tensor_dict(tensor_dict, with_opt_vars=False, device='cpu')

Set the tensor dictionary.

Parameters:
  • tensor_dict – The tensor dictionary

  • with_opt_vars (bool) – Return the tensor dictionary including the optimizer tensors (Default=False).