Exchanger refactor #33
Conversation
| def __init__(self, parameter_packer: ParameterPacker[T]) -> None: | ||
| super().__init__() | ||
| self.parameter_packer = parameter_packer | ||
|
|
There was a problem hiding this comment.
By default, ParameterExchangerWithPacking exchanges all parameters like FullParameterExchanger. We can create subclasses of it if we wish to use different exchanging methods. e.g.: NormDriftLayerExchanger
| class NormDriftLayerExchanger(ParameterExchangerWithPacking[List[str]]): | ||
| def __init__(self, threshold: Scalar) -> None: | ||
| """ | ||
| self.initial_model represents each client's local model at the beginning of each round of training. |
There was a problem hiding this comment.
I would suggest editing this comment a bit so that it better reflects the new structure of the class.
| initial_model_states = self.initial_model.state_dict() | ||
| initial_model_states = initial_model.state_dict() | ||
| model_states = model.state_dict() | ||
| for layer_name in model_states: |
There was a problem hiding this comment.
I realized we forgot to reformat this for loop in the previous PR, so I'd suggest we do it here.
for layer_name, layer_param in model_states.items():
Then you don't need to do the extra indexing step in line 58 below.
emersodb
left a comment
emersodb
left a comment
There was a problem hiding this comment.
Two small comments from me. Otherwise, this is a really great change in my opinion.
|
Good catch! I have fixed the comments and reformatted the for loop. |
jewelltaylor
left a comment
jewelltaylor
left a comment
There was a problem hiding this comment.
Everything looks good from my end! Decoupling the Parameter Packer from the exchanger and adding generic typing makes the code much cleaner. I will leave it to David for the final approval to make sure he is happy with how you addressed his suggestions (which seems to be the case but not for me to say)
PR Type
[Feature | Fix | Documentation | Other() ]
Feature
Short Description
The main goal of this PR is to redesign the parameter packing component of the parameter exchanger. The packing component is now its own module
parameter_packer, and is decoupled from the exchanger itself. I used Generics to support better typing for the packer class. I also used the packers to refactor the strategies.Tests Added
N/A