How to reset Model parameters to random values after training it?

I have a requirement where I want to reset my model parameters to random values after a constraint is satisfied.

class SamplerModel(nn.Module):
    def __init__(self, input_size, hidden_size, n_dims):
        super(SamplerModel, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc_mu = nn.Linear(hidden_size, 1)  # Output layer for mu
        self.fc_sigma = nn.Linear(hidden_size, n_dims)  # Output layer for sigma

    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        mu = self.fc_mu(x)
        sigma = torch.exp(self.fc_sigma(x))  # Apply exponential to ensure sigma is positive
        return mu, sigma

Above is my model. I want to reset the values after certain constraint satisfaction. how do I do it.

You can call the .reset_parameters() method of each module containing trainable parameters or you could write your custom initialization method and call it via model.apply.

Is there a way to do this with nn.Sequential models?

Yes, you can iterate the submodules:

model = nn.Sequential(
    nn.Linear(1, 1),
    nn.ReLU(),
    nn.Linear(1, 1)
)

print(model[0].weight)
# Parameter containing:
# tensor([[-0.8953]], requires_grad=True)

for m in model.modules():
    if hasattr(m, "reset_parameters"):
        m.reset_parameters()

print(model[0].weight)
# Parameter containing:
# tensor([[0.9524]], requires_grad=True)
1 Like