Optimize dynamically generated modules

[ajoino]

I have a network that adds modules dynamically depending on the data, but since those modules are added after configure_optimizers is run they will not be optimized during training. I have read that you can customize the manual_backward to do something similar to this but I am not sure exactly how.

In more detail, my network basically looks like this:

class DynNetwork(pl.LightningModule):
    def __init__(self, dim=128):
        self.dynamic_module = DynamicModule(dim) # This is the network that contains the dynamic modules, all added with self.add_module(...)
        self.other_module = nn.Linear(dim, 1)

    def configure_optimizers():
        return torch.optim.Adam(self.parameters())

How would I add a manual_backward here, or could this be done with some trick to make the optimizer evaluate self.parameters() each run?

[ananthsub]

Out of curiosity, are you using PyTorch lazy modules?

Is it possible to initialize your module in the setup hook?

Or if you have an example batch of data before training, is it possible to initialize your module outside of the lightning module and then pass into the lightning module?

[ajoino]

I'm not using lazy modules no, and I cannot initialize the dynamically generated modules in the setup hook because the modules are created when it sees a new "data element".

I can get around my problem by first running model.predict on a subset of the training data as you suggest, but I'm hoping there's a better way.

[ananthsub]

@ajoino there are 2 aspects:

• You could control the optimization using https://pytorch-lightning.readthedocs.io/en/latest/common/optimizers.html#manual-optimization so the sequence of loss computation, backwards, and optimizer steps are entirely under your control
• However, what's unclear to me is how the trainer will know about the optimizers you create on the fly in order to checkpoint their states
• What's also unclear to me is how this would work in a distributed setting

[tchaton]

Dear @ajoino,

We currently have similar mechanism within the BaseFinetuning Callback.
You could add a new param group within the existing optimizers once your new layers are created.
Would this solve your problem ?

Best,
T.C