Setting some elements of layer parameters to 0 for each training iteration

Hi everyone,

I want to fine-tune a pre-trained network with same data as it was trained.
I need help in setting some element of conv2d layer weight (and bias too) to 0 for each training iteration and others are trained as normal. My goal is to have 0 outputs where set as 0.

For example, my layer has shape
conv1.weight.shape # torch.Size([32, 64, 7, 7])
and I have a list of unit that need to be set to 0 zeroed_units = [1, 3, 5, 7, 9]
and I would like it to have the tensor of conv1.weight at index [1, 3, 5, 7, 9] = 0, i.e. torch([32,zeroed_units,7,7] = 0)

  • I have thought about setting these units 0 and set requires_grad=False. However, according to google, requires_grad only allow for entire tensor layer not part of it.
  • Another way is to apply a non-trainable mask buffer of 1’s and set 0 where needed after this conv2d and does not touch the conv2d layer. This method could give me the desired output. However, when training in this way, the conv2d layer weight distribution would have the same distribution as before. The number might change but the distribution is not changed much. I would like to see the effect of different weight’s distribution on the final results when setting units to 0.

Thank you.

With a bit of luck, I found some answers that are close to what I want to do.

Can.detach() work for parts of the layer weights?
and
Update only sub-elements of weights

It seems that these are not the optimal solution in terms of saving resources since it uses the gradient mask and sets the gradient to 0 where needed. However, this seems to be the only possible solution for me right now.