Allow subset of layer weights to be updated

I have a network and after a few conv layers, it has 3 heads: self.fc_pi, self.fc_q and self.fc_beta. I have frozen all of the conv layer weights and am allowing updates only to self.fc_q and self.fc_beta by adding only their weights to the optimizer as:

self.optimizer = config.optimizer_fn(

The self.fc_pi is a linear layer that connects 200 neurons to 16 neurons so when I do print(list([0].shape), I get [16,200] and when I do print(list([1].shape), I get [16], both of which respectively represent the weights and biases of this linear layer.

I would like to add a subset of these weights and biases to the optimizer so that only the weights connected to the first 4 (out of the 16 neurons), and their biases, are updated. I have tried this:

self.optimizer.add_param_group({'params': list([0][0:4]})
self.optimizer.add_param_group({'params': list([1][0:4]})

but this raises the error:

ValueError: can't optimize a non-leaf Tensor

How can I correctly add a subset of the layer weights and biases to the optimizer so that only a subset of weights is updated and not all the weights in the layer.

After the backward call, and before the optimizier update, you can access the gradient tensor (param.grad) and assign zero to the parts that you want to freeze/not update. This way these weights are left unchanged after the optim.step call.

1 Like