How to train only a certain number of parameters of a PyTorch tensor

1

Hello,

0

I have a 2 X 2 tensor in PyTorch. I want the first row entries of this tensor (matrix) to be fixed, i.e., I do not want the first row entries to be parameters that are trained using gradient descent, and the second row entries be updated by the training process (using gradient descent). I saw online that I can use the register_buffer function in PyTorch to create a gradient_mask but I am not sure about the exact implementation.

matrix = torch.nn.Parameter(torch.zeros(2, 2))
register_buffer(‘gradient_mask’, ???)

Also, after I create this gradient_mask I do not know how to apply it to my variable: matrix. Any thoughts? Thanks.

2

I believe you can manually store the first row and then restore it after training - I created a simple example

import torch
import torch.optim as optim
matrix = torch.nn.Parameter(torch.randn(2, 2))
print("Original Parameters",matrix )
gradient_mask = torch.tensor([[0, 0], [1, 1]], dtype=torch.float32)
optimizer = optim.SGD([matrix], lr=0.01)
input_data = torch.randn(2, 1)
output = torch.matmul(matrix, input_data)

#Calculate the loss
loss = output.sum()
optimizer.zero_grad()
loss.backward()

# Apply the gradient mask to zero out gradients for the first row
with torch.no_grad():
    matrix.grad *= gradient_mask

optimizer.step()

# Restore the first row to its initial value
initial_first_row = matrix[0].clone()
with torch.no_grad():
    matrix[0] = initial_first_row
print("Parameters after restoring the first row:\n", matrix)
1 Like