Train one optimizer/ model each batch

[sachinruk]

I am wondering what the default behaviour of lightning is when it comes to training multiple optimizers. I am not entirely sure, but it feels like maybe lightning is optimizing all available parameters from all available models in each backward_step.

Here is the desired behaviour:

        optimizers = self.configure_optimizers()
        models = [self.query_model, self.doc_model, self.position_model]
        
        j = 0
        optimizer = optimizers[j]
        model = models[j]
        for _ in range(epochs):
            for i, (x, y) in enumerate(train_dl):
                loss = self.training_step(x, y, j)
                loss = loss / self.accumulation_steps
                loss.backward()
                if (i + 1) % self.accumulation_steps == 0:
                    optimizer.step()
                    model.zero_grad()
                    j = (j + 1) % len(optimizers)
                    optimizer = optimizers[j]
                    model = models[j]

I am running pl version 1.1.8. Happy to upgrade if it helps.

[carmocca]

You can check out the automatic_optimization pseudocode for multiple optimizers here:
https://pytorch-lightning.readthedocs.io/en/latest/common/optimizers.html#automatic-optimization

We also make sure only the gradients of the current optimizer’s parameters are calculated in the training step to prevent dangling gradients in multiple-optimizer setups

[sachinruk]

Hey. Thanks for that, but even though it's similar to what I want to do, it's not quite there. Looking at that example which I've copied below, it seems like it would do all 3 optimizers (in my case) in each batch. Is there a way to do different batches for each optimizer? Or for that matter, ideally I want to do X number of batches before switching optimizers.

for epoch in epochs:
    for batch in data:
        for opt in optimizers:
            loss = model.training_step(batch, batch_idx, optimizer_idx)
            opt.zero_grad()
            loss.backward()
            opt.step()

[tchaton]

Dear @sachinruk, here is a slightly more accurate.

for epoch in epochs:
    for batch in data:
        for opt in optimizers:
            opt.toggle_parameters(model)
            loss = model.training_step(batch, batch_idx, optimizer_idx)
            opt.zero_grad()
            loss.backward()
            opt.step()
            opt.untoggle_parameters(model)

Here is a test with 3 optimizers to better understand how it works :)

As you can see, when running optimizer_0 which contains layer 1 and layer 2, layer 3 parameters are frozen. However, layer 1 and layer 2 still preserve the frozen parameters set by the users.

This rotate when passing to optimizer 1, and optimizer_2, etc ...

def test_toggle_untoggle_3_optimizers_shared_parameters(tmpdir):

    class TestModel(BoringModel):

        def __init__(self):
            super().__init__()
            self.layer_1 = nn.Sequential(
                nn.Linear(32, 32),
                nn.ReLU(),
                nn.Linear(32, 32),
                nn.ReLU(),
                nn.Linear(32, 32),
            )

            self.layer_2 = nn.Sequential(
                nn.ReLU(),
                nn.Linear(32, 32),
                nn.ReLU(),
                nn.Linear(32, 32),
                nn.ReLU(),
                nn.Linear(32, 2),
            )

            self.layer_3 = nn.Sequential(
                nn.ReLU(),
                nn.Linear(32, 32),
                nn.ReLU(),
                nn.Linear(32, 32),
                nn.ReLU(),
                nn.Linear(32, 2),
            )

            # set some weights to False to check untoggle works as expected.
            self.layer_1[2].weight.requires_grad = False
            self.layer_1[4].weight.requires_grad = False

            self.layer_2[1].weight.requires_grad = False
            self.layer_2[3].weight.requires_grad = False

            self.layer_3[1].weight.requires_grad = False
            self.layer_3[5].weight.requires_grad = False

        def optimizer_step(
            self,
            current_epoch,
            batch_nb,
            optimizer,
            optimizer_idx,
            closure,
            on_tpu=False,
            using_native_amp=False,
            using_lbfgs=False
        ):
            if optimizer_idx == 0:
                assert self.layer_1[0].weight.requires_grad is True
                assert self.layer_1[2].weight.requires_grad is False
                assert self.layer_1[4].weight.requires_grad is False

                assert self.layer_2[1].weight.requires_grad is False
                assert self.layer_2[3].weight.requires_grad is False
                assert self.layer_2[5].weight.requires_grad is True

                assert self.layer_3[1].weight.requires_grad is False
                assert self.layer_3[3].weight.requires_grad is False
                assert self.layer_3[5].weight.requires_grad is False

            if optimizer_idx == 1:
                assert self.layer_1[0].weight.requires_grad is False
                assert self.layer_1[2].weight.requires_grad is False
                assert self.layer_1[4].weight.requires_grad is False

                assert self.layer_2[1].weight.requires_grad is False
                assert self.layer_2[3].weight.requires_grad is False
                assert self.layer_2[5].weight.requires_grad is True

                assert self.layer_3[1].weight.requires_grad is False
                assert self.layer_3[3].weight.requires_grad is True
                assert self.layer_3[5].weight.requires_grad is False

            if optimizer_idx == 2:
                assert self.layer_1[0].weight.requires_grad is True
                assert self.layer_1[2].weight.requires_grad is False
                assert self.layer_1[4].weight.requires_grad is False

                assert self.layer_2[1].weight.requires_grad is False
                assert self.layer_2[3].weight.requires_grad is False
                assert self.layer_2[5].weight.requires_grad is False

                assert self.layer_3[1].weight.requires_grad is False
                assert self.layer_3[3].weight.requires_grad is True
                assert self.layer_3[5].weight.requires_grad is False

            optimizer.step(closure=closure)

        def training_step(self, batch, batch_idx, optimizer_idx=None):
            loss = super().training_step(batch, batch_idx)
            # make sure the model is untoggle when returning None
            return loss if batch_idx % 2 == 0 else None

        @staticmethod
        def combine_generators(gen_1, gen_2):
            for p in gen_1:
                yield p
            for p in gen_2:
                yield p

        def configure_optimizers(self):
            optimizer_1 = SGD(self.combine_generators(
                self.layer_1.parameters(),
                self.layer_2.parameters(),
            ), lr=0.1)
            optimizer_2 = Adam(self.combine_generators(
                self.layer_2.parameters(),
                self.layer_3.parameters(),
            ), lr=0.1)
            optimizer_3 = SGD(self.combine_generators(
                self.layer_3.parameters(),
                self.layer_1.parameters(),
            ), lr=0.1)
            return [optimizer_1, optimizer_2, optimizer_3]

    model = TestModel()
    model.training_epoch_end = None

    trainer = Trainer(
        max_epochs=1,
        default_root_dir=tmpdir,
        limit_train_batches=8,
        accumulate_grad_batches=2,
    )

    trainer.fit(model)

Find the test there: https://github.com/PyTorchLightning/pytorch-lightning/blob/ced2c94a3eceecfde2a9a2eac7a3897530c373fb/tests/core/test_lightning_module.py#L206

Best,
T.C