how to set the gradient for a network in pytorch

Question

I have a model in pytorch. The model can take any shape but lets assume this is the model

torch_model =  Sequential(
    Flatten(),
    Linear(28 * 28, 256),
    Dropout(.4),
    ReLU(),
    BatchNorm1d(256),
    ReLU(),
    Linear(256, 128),
    Dropout(.4),
    ReLU(),
    BatchNorm1d(128),
    ReLU(),
    Linear(128, 10),
    Softmax()
)

I am using SGD optimizer, I want to set the gradient for each of the layers so the SGD algorithm will move the parameters in the direction I want.

Lets say I want all the gradients for all the layers to be ones (torch.ones_like(gradient_shape)) how can I do this? Thanks?

Answer 1

In PyTorch, with a model defined as yours above, you can iterate over the layers like this:

for layer in list(torch_model.modules())[1:]:
  print(layer)

You have to add the [1:] since the first module returned is the sequential module itself. In any layer, you can access the weights with layer.weight. However, it is important to remember that some layers, like Flatten and Dropout, don't have weights. A way to check, and then add 1 to each weight would be:

for layer in list(torch_model.modules())[1:]:
  if hasattr(layer, 'weight'):
    with torch.no_grad():
      for i in range(layer.weight.shape[0]):
          layer.weight[i] = layer.weight[i] + 1

I tested the above on your model and it does add 1 to every weight. Worth noting that it won't work without torch.no_grad() as you don't want pytorch tracking the changes.