[Python] Freeze then unfreeze gradients of a subset of tensor in PyTorch, using...

I'm a self-taught PyTorch user getting stuck on a possibly-trivial-or-not question; it has been confusing me and all my AI tools for months!

Let us say I have a tensor X of shape [N, ] (possibly a gradient vector) and I wish to freeze part of it during a part of my script (which can be indicated by a mask of shape [N, ] of boolean values, where False indicates I wish to freeze the gradient of the corresponding element in X, and True indicates I wish to unfreeze the element and allows its gradient to back-propagate. After part 1 of my code is run, I will run part 2 where I'll simply reverse mask as ~mask; so all trained elements in X will be frozen and all previously frozen elements in X will be unfrozen and resume training.

I have tried two different approaches: 1) use the register_hook() function to assign a masking function; I have also tried either explicitly defining the value of mask as either a dynamic or static object in the memory; 2) call with torch.no_grad() followed by X.grad()[mask] = 0.

In both scenarios the first part of my code works perfectly; for both approach 1) and 2) I can correctly freeze the elements of X that I wish to fix while training the others. The problem comes when I tried to unfreeze the rest while freezing the trained ones. In both approach this phenomenon happens:

• 
  if I set all parameters to "unfreeze", no matter via approach 1) or 2), all parameters will unfreeze as expected;
  
  
• 
  however, if I tried to "freeze" an element on the second attempt either via register_hook() with a changed mask; via .remove() my old hook then applying a new one; or via setting .grad() = 0 for the corresponding element, the element will NOT be frozen.
  

In summary, once I manipulate the gradient of my tensor X for the first time it works well as intended; as soon as I did this for a second time, I can only unfreeze it, not freeze.

I have consulted with various sources on the internet; the PyTorch documentation suggests any manipulation done with the hook should take place after a .backward(); and AI tools do not help as they are heavily hallucinated, probably meaning no one has encountered this problem before.

Any help would be greatly appreciated!!

Continue reading...