Can't pickle SWA: AttributeError: 'SWA' object has no attribute 'defaults'

[vitusbenson]

For some reason the self.defaults is not properly set when in manual mode so it can not be pickled. Workaround: Change self.optimizer.defaults to optimizer.defaults and self.optimizer.param_groups to optimizer.param_groups and self.optimizer.state to optimizer.state.

https://github.com/pytorch/contrib/blob/c545fedf4f73c8e95f91fd81f2d5bf7fa9c62a61/torchcontrib/optim/swa.py#L122

[rensjochemsen]

I encountered the same error and tried to fix it your way. This is what I had:

opt_maker = lambda x: SWA(optim.AdamW(x.parameters(), lr=self.learning_rate, eps=1e-4)) if self.swa \
                                else optim.AdamW(x.parameters(), lr=self.learning_rate, eps=1e-4)
self.optimizer = opt_maker(self.network)

After fix:

opt_maker = lambda x: SWA(optim.AdamW(x.parameters(), lr=self.learning_rate, eps=1e-4)) if self.swa \
                                else optim.AdamW(x.parameters(), lr=self.learning_rate, eps=1e-4)
self.optimizer = opt_maker(self.network)
self.optimizer.param_groups = self.optimizer.optimizer.param_groups
self.optimizer.state = self.optimizer.optimizer.state
self.optimizer.defaults=self.optimizer.optimizer.defaults

Maybe this was not what you meant. For me it resulted in the the error:

BrokenPipeError: [Errno 32] Broken pipe