Support for transfer learning with feature extraction

[gconst02]

Is your feature request related to a problem? Please describe.
Current version of federated_client.py optimizes all parameters of the model.

self.optimizer = optimizer(model.parameters(), **optimizer_args)

For transfer learning with feature extraction, only the classifier needs to be optimized.

Describe the solution you'd like
train_config.py already accepts optimizer_args but cannot specify params parameter of the optimizer (if you do default you get TypeError: __init__() got multiple values for argument 'params' due to the default model.parameters()).

This can be modified to optionally set params if appears in optimizer_args or else use current behavior , i.e.,

optimizer = getattr(th.optim, optimizer_name)
optimizer_args.setdefault("params", model.parameters())
self.optimizer = optimizer(**optimizer_args)

Describe alternatives you've considered
None

Additional context
None

[AniTho]

Is this problem somewhat like what we do with fastai library that is at the beginning, we freeze all the layer and only train and optimize the final classifier layer and then we unfreeze the top layers to start optimizing the complete model?

[gconst02]

@AniTho I haven't used fastai but from what you described:

we freeze all the layer and only train and optimize the final classifier layer

-> Is transfer learning with feature extraction

and then we unfreeze the top layers to start optimizing the complete model

-> Is transfer learning with fine-tuning

With this change, we will have the flexibility to specify which params to optimize.

[github-actions]

This issue has been marked stale because it has been open 30 days with no activity. Leave a comment or remove the stale label to unmark it. Otherwise, this will be closed in 7 days.