Attaching an optimizer as an attribute to a nn.Module instance

[shartzog (Sam Hartzog)]

@Giobio
What’s the driver for storing the optimizers within your nn.Module class? I don’t think that would break anything, but it seems unrelated to the ability to use multiple optimizers. For example, you could probably handle it with a custom optimizer class as described here:

Two optimizers for one model

Use the basic knowledge of software engineering. class MultipleOptimizer(object): def __init__(*op): self.optimizers = op def zero_grad(self): for op in self.optimizers: op.zero_grad() def step(self): for op in self.optimizers: op.step() opt = MultipleOptimizer(optimizer1(params1, lr=lr1), optimizer2(params2, lr=lr2)) loss.backward() opt.zero_grad() opt.step()

There’s some other good in that post as well. Hope that helps!
–SEH