Sideffects when changind some elements of nn.Parameter
I want to have some nn.Parameter that depend on data, e.g. different values for different people. I realized while training one person, the results of the other people also change. I created a small demo below.
I tried to make everything deterministic and as expected, changing the size of nn.Parameter does not change the results. After training for 100 epochs, I always get [1.0605, 1.0715, 1.0622, 0.9702] for all_calibration_params[0]. But all other parameters of all_calibration_params also change, and they all change by the same value, with all_calibration_params = nn.Parameter(torch.ones(10, 4)) I get [0.8455, 0.8455, 0.8455, 0.8455] for all 9 other values. I would expect these values to stay [1, 1, 1, 1]. Can anyone please explain to me, why the other values change even though I do not use them during training?
Thank you.
import torch
from pytorch_lightning import LightningModule, Trainer, seed_everything
from torch import nn
from torch.utils.data import Dataset, DataLoader
import torch.nn.functional as F


class DemoDataset(Dataset):
    &quot;&quot;&quot;
    Generate random data, only for person 0.
    Data is always random but output is always 1.
    &quot;&quot;&quot;

    def __init__(self):
        pass

    def __len__(self) -&gt; int:
        return 1_000

    def __getitem__(self, idx):
        return {
            'person': torch.zeros(1),
            'data': torch.rand(1),
            'output': torch.ones(1),
        }


class DemoModel(LightningModule):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.all_calibration_params = nn.Parameter(torch.ones(3, 4))  # all values are 1

        self.fc_1 = nn.Linear(1, 4)
        self.fc_2 = nn.Linear(4 + 4, 1)

    def forward(self, person: torch.Tensor, data: torch.Tensor):
        x = self.fc_1(data)
        calibration_params = self.all_calibration_params[person].squeeze(1)
        x = torch.cat((x, calibration_params), dim=1)
        x = self.fc_2(x)

        return x


class Model(DemoModel):
    def __init__(self, learning_rate: float, weight_decay: float, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.learning_rate = learning_rate
        self.weight_decay = weight_decay

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)

    def training_step(self, train_batch: dict, batch_idx: int):
        data, person, output = train_batch['data'], train_batch['person'].long(), train_batch['output']

        pred = self(person, data)

        loss = F.mse_loss(pred, output)
        return loss


if __name__ == '__main__':
    seed_everything(42)

    model = Model(learning_rate=1e-4, weight_decay=1e-5)

    dataset = DemoDataset()
    dataloader = DataLoader(dataset, batch_size=64)

    print(model.all_calibration_params)

    trainer = Trainer(
        max_epochs=100,
        deterministic=True
    )
    trainer.fit(model, dataloader)

    print(model.all_calibration_params)

    # model.all_calibration_params[1] should be [1, 1, 1, 1]
    print(model.all_calibration_params[1], torch.ones(1, 4))

Output:
tensor([[1.0605, 1.0715, 1.0622, 0.9702],
        [0.8455, 0.8455, 0.8455, 0.8455],
        [0.8455, 0.8455, 0.8455, 0.8455]], requires_grad=True)

EDIT:
I added a small test if the other elements in nn.Parameter change the output, but the do not, so I do not understand why they are affected while training.
batch = next(iter(dataloader))
data, person = batch['data'], batch['person'].long()
output_org = model(person, data)
print(output_org[0])

model.freeze()
model.all_calibration_params[1:] = torch.ones(2, 4)
model.unfreeze()

output_mod = model(person, data)
print(output_mod[0])
print(torch.equal(output_org, output_mod))

Output:
tensor([0.6334], grad_fn=&lt;SelectBackward&gt;)
tensor([0.6334], grad_fn=&lt;SelectBackward&gt;)
True


Your dataset always returns a 0 index for the 'person', i.e. self.all_calibration_params[person] will always select the first rows.
The reason the other two rows are getting updated is because of the weight decay that was set on the optimizer. This will update all parameters regardless of whether they were used or not during the model inference.
