I am facing a weird issue when trying to optimize parameters in a custom module. In particular, the weights of certain parameters are not updated, even if their gradients are computed.
My CustomModule
is defined as:

A list of Batch Normalization (BN) layers, in which, only one is chosen for the forward/backward pass. That is, the model will learn different gammas/betas (based on the selected BN layer).

A fullyconnected layer.
Of course, this structure, defined as such, does not make sense. I just wanted to focus on the occurred issue. Below, the structure of the CustomModule
.
class CustomModule(nn.Module):
def __init__(self, num_batch_norm):
super(CustomModule, self).__init__()
self.batch_norm_list = [nn.BatchNorm1d(10) for _ in range(num_batch_norm)]
self.fc = nn.Linear(10, 2)
def forward(self, x, num_bn):
x = self.batch_norm_list[num_bn](x)
x = self.fc(x)
return x
I create a model
(instance of CustomModule
) with three
BN layers and SGD
optimizer. After, I define test variables x
and y
to perform a single pass through the model
(including the backward pass).
torch.manual_seed(0)
# Create a 'CustomModule' with three BN layers.
model = CustomModule(3)
optimizer = torch.optim.SGD(model.parameters(), lr=1e2)
# Define test variables: the input (x) and the target (y).
x = torch.rand(5, 10)
y = torch.rand(5, 2)
# Pass 'x' through the 2nd (index=1) BN layer.
out = model(x, 1)
loss = torch.linalg.norm(out  y)
model.zero_grad()
loss.backward()
optimizer.step()
I check that the gradients for the 2nd BN layer (chosen in the forward pass) are computed.
# Check the weights' gradients in the 2nd BN layer.
model.batch_norm_list[1].weight.grad
# tensor([ 0.0810, 0.0052, 0.2871, 0.3091, 0.2437, 0.0916, 0.1410, 0.4406,
# 0.0288, 0.0740])
However, when I check the weights (gammas) after optimization, I noticed that they are not updated (in batch normalization, gammas (weights) are initialized by one
).
# Check the weights in the 2nd BN layer (after optimization)
model.batch_norm_list[1].weight
# Parameter containing:
# tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.], requires_grad=True)
Having the initial weights (a tensor of ones) and gradients, I expected the updated weights to be like below (using the SGD
optimizer).
# Expected values
weights = model.batch_norm_list[1].weight
gradients = model.batch_norm_list[1].weight.grad
lr = 1e2
expected_weights = weights  lr * gradients
# tensor([0.9992, 0.9999, 0.9971, 0.9969, 0.9976, 0.9991, 0.9986, 0.9956, 1.0003,
# 0.9993], grad_fn=<SubBackward0>)
So, I want to know why this issue encounters and how to solve it.