Transfer Learning

apax.transfer_learning.parameter_transfer.black_list_param_transfer(source_params: FrozenDict | dict, target_params: FrozenDict | dict, param_black_list: list[str]) → FrozenDict

Transfer parameters from one dictionary to another, while keeping

some key-value pairs unchanged.

Parameters:

• source_params (Union[FrozenDict, dict]) – source parameters

• target_params (Union[FrozenDict, dict]) – target parameters

• param_black_list (list[str]) – list of keys to keep unchanged.

Returns:

target_params with key-value pairs updated.

Return type:

transfered_target (dict)

apax.transfer_learning.parameter_transfer.transfer_parameters(state: TrainState, ckpt_config: TransferLearningConfig) → TrainState

Transfer the parameters from the checkpoint to the train state.

Parameters:

• state (TrainState) – train state

• ckpt_config (TransferLearningConfig) – transfer learning configuration

Returns:

TrainState with the params attribute updated

according to the transfer learning configuration.

Return type:

state (TrainState)