How to freeze part of the weight matrix during training?

Asgard · November 20, 2021, 7:24pm

I have:

model = Sequential(Linear(8,5), Linear(5,3), Linear(3,1))

Now, for model[1].weight, I want to freeze part of this weight matrix during the training procedure i.e.,

model[1].weight[0:3,0:3]

values shouldn’t vary during the backpropagation. Apparently, model[1].weight[0:3,0:3].detach() isn’t working as the weights vary during the training. And model[1].weight[0:3,0:3].requires_grad=False throws error.

So, how to achieve this? Thank you in advance.

InnovArul · November 20, 2021, 9:17pm

I am not sure if any predefined way exists in pytorch.
You may need to write a custom linear layer to achieve it, something like below:

gist.github.com

https://gist.github.com/InnovArul/ee2986feff215fed94401ec83367408b

linear_partial_freeze.py

import torch, torch.nn as nn                                                         
import torch.optim as optim, torch.nn.functional as F                                
                                                                                     
class CustomLinear(nn.Module):                                                       
    def __init__(self, mask):                                                        
        super().__init__()                                                           
        self.register_buffer("mask", mask)                                           
        out_channels, in_channels = mask.shape                                       
        self.weight = nn.Parameter(torch.randn(out_channels, in_channels))

This file has been truncated. show original