Stop gradient computation for specifc weight in wt matrix

username: VIVEK_RUHELA
cooked: <p>I am trying to develop a model for parameter estimation in which some parameters are allowed to train and some parameters are not allowed to train. e.g. In following matrix <code>w_matrix = np.array([[w1,0.0,0.0,-w2],[w3,-w4,0.0,0.0],[0.0,w5,-w6,0.0]]</code> all <code>0.0</code> are fixed weight and all <code>w</code> are allowed to update.</p>
<p>My attempt is as follows:</p>
<pre><code class="lang-auto">class NeuralNet(nn.Module):
    
    def __init__(self,param):
        super(NeuralNet, self).__init__() 
        param.append(0.0)
        self.gamma = torch.tensor(param[0])
        self.gamma = Variable(self.gamma,
                              requires_grad=True).type(dtype)
        self.alpha_xy = torch.tensor(param[1])
        self.aplha_xy = Variable(self.alpha_xy,
                                 requires_grad=True).type(dtype)
        self.beta_y = torch.tensor(param[2])
        self.beta_y = Variable(self.beta_y,
                               requires_grad=True).type(dtype)
        self.alpha0 = torch.tensor(param[3])
        self.alpha0 = Variable(self.alpha0,
                               requires_grad=True).type(dtype)
        self.alpha_y = torch.tensor(param[4])
        self.alpha_y = Variable(self.alpha_y,
                               requires_grad=True).type(dtype)
        self.alpha1 = torch.tensor(param[5])
        self.alpha1 = Variable(self.alpha1,
                               requires_grad=True).type(dtype)
        self.alpha2 = torch.tensor(param[6])
        self.alpha2 = Variable(self.alpha2,
                               requires_grad=True).type(dtype)
        self.alpha3 = torch.tensor(param[7])
        self.alpha3 = Variable(self.alpha3,
                               requires_grad=True).type(dtype)
        self.zero_fix_wt = torch.tensor(param[8])
        self.zero_fix_wt = Variable(self.zero_fix_wt,
                               requires_grad=False).type(dtype)
        
        # Production Matrix
        self.w_production = torch.tensor([[self.gamma, self.zero_fix_wt, self.zero_fix_wt, self.zero_fix_wt],
                                         [self.beta_y, self.zero_fix_wt, self.zero_fix_wt self.zero_fix_wt],
                                         [self.zero_fix_wt, self.alpha0, self.zero_fix_wt, self.zero_fix_wt]])
        self.w_production = Variable(self.w_production.to(device)).type(dtype)
        self.w_production = nn.Parameter(self.w_production)
        
        # Degradation Matrix
        self.w_decay = torch.tensor([[-self.alpha1, self.zero_fix_wt, self.zero_fix_wt, self.zero_fix_wt],
                                     [self.zero_fix_wt, -self.alpha2, self.zero_fix_wt, self.zero_fix_wt],
                                     [self.zero_fix_wt, self.zero_fix_wt, -self.alpha_y, self.zero_fix_wt]])
        self.w_decay = Variable(self.w_decay.to(device)).type(dtype)
        self.w_decay = nn.Parameter(self.w_decay)        
        
        # Cross-Talk Matrix
        self.w_cross_talk = torch.tensor([[self.zero_fix_wt, self.zero_fix_wt, self.zero_fix_wt, -self.alpha_xy],
                                         [self.zero_fix_wt, self.zero_fix_wt, self.zero_fix_wt, self.zero_fix_wt],
                                         [self.zero_fix_wt, self.zero_fix_wt, self.zero_fix_wt, -self.alpha3]])
        self.w_cross_talk = Variable(self.w_cross_talk.to(device)).type(dtype)
        self.w_cross_talk = nn.Parameter(self.w_cross_talk)
        
        
        
    def forward(self,input): 
        xy_term = input[0][0] * input[0][2]        
        input_new = torch.cat((input,xy_term.view([-1,1])),1)
        input_new = input_new.to(device,non_blocking=True)        
        hidden_state = (self.w_production + self.w_decay + self.w_cross_talk).mm(torch.transpose(input_new,0,1))
        out = hidden_state.view([-1,3]) + input
        return (out, hidden_state)
    
   

# parameter initialization
init_params = [2.0,3.7,1.5,1.1,0.9,0.1,0.9,0.01]

# define model in GPU
model = NeuralNet(init_params)
model = model.to(device)

loss_function = nn.MSELoss()

optimizer = torch.optim.SGD(model.parameters(), lr=1e-3,momentum = 0.9)

start = datetime.now()
loss_1 = []

for i in tqdm(range(epochs)):
    total_loss = 0
    
    for j in range(x_train.size(0)):
        optimizer.zero_grad()
        input = x_train[j:(j+1)]
        target = y_train[j:(j+1)]        
        
        input = input.to(device,non_blocking=True)
        target = target.to(device,non_blocking=True)        
        
        (pred, _) = model(input)
        loss = loss_function(pred,target)
        total_loss += loss        
        loss.backward()   
        optimizer.step()  
        print(model.w_production)
                    
                   
    
    # Early Stopping
    loss_1.append(total_loss.item())        
    if len(loss_1) &gt; 1 and np.abs(loss_1[-2] - total_loss.item()) &lt;= 0.10:
        print('Early Stopping due to dIfference in error is less than threshold (0.10)')
        print('Stopping training at %d epochs with total loss = %d' %(i,total_loss))
        break


    if i % 50 == 0:
        print("Epoch: {} loss {}".format(i, total_loss))
        
end = datetime.now()
time_taken = end - start
print('Execution Time: ',time_taken)
</code></pre>
<p>Here problem I am facing is that <code>loss.backward()</code> and <code>optimizer.step()</code> compute the gradient and adjust it for all weight even I have set <code>requires_grad=False</code> for some weights. This is happening because of <code>requires_grad=False</code> workes either on the complete matrix or not at all.</p>
<p>My question is how to tell model not to touch fixed weight and never compute the gradient for the specifically defined fixed weight for which <code>requires_grad=False</code> is mentioned.</p>

username: vainaijr
cooked: <p>I think if you register hook and zero out gradients for mask, then it will work, for example,<br>
‘’’<br>
def hook_fn(grad):<br>
out = grad.clone()<br>
out[mask] = 0<br>
return out</p>
<p>layer.weight.register_hook(hook_fn)<br>
‘’’</p>

username: VIVEK_RUHELA
cooked: <p>Thanks. Isn’t is similar to manually zeroing the gradient of selected weight. If this is so, I don’t think that this will work because 0.0 in weight matrix represent that link is not present but for model link is present with zero weight. These two are very much different aspect, one is constraint optimization aspect and another is model architectural aspect.</p>

username: VIVEK_RUHELA
cooked: <p>Upto my guess, this problem can be solved by the trick that</p>
<p><strong>Stop <code>loss.backward()</code> to compute error at each computational nodes and implement .backward() only at permitted nodes so that <code>optimizer.step()</code> update only permitted weights and don’t update fixed weight</strong></p>
<p>But now the question is  how to customize <code>loss.backward()</code> or is there any other efficient and fast way?</p>

username: VIVEK_RUHELA
cooked: <p>My hypothesis of writing customized loss function to deal with customized model was wrong Then I tried something different.<br>
I have figured out one solution using this <a href="https://github.com/uchida-takumi/CustomizedLinear" rel="noopener nofollow ugc">reference</a> (but again there are some issues in the model I am not able to understand why this is happening)</p>
<p>My codes are as follows:</p>
<pre><code class="lang-auto">class CustomizedForwardBackwardFunc(torch.autograd.Function):
    '''
        Autograd function which masks it's weight by 'mask'
    '''
    
    @staticmethod
    def forward(ctx, input, production_weight, decay_weight, 
                cross_talk_weight, bias, prod_mask, decay_mask, cross_mask):
        
        if prod_mask.size() == production_weight.size():
            production_weight = torch.mul(production_weight, prod_mask)
            decay_weight = torch.mul(decay_weight, decay_mask)
            cross_talk_weight = torch.mul(cross_talk_weight, cross_mask)
        else:            
            raise Exception('Please check your production mask matrix size')
            
        xy_term = input[0][0] * input[0][2]        
        input_new = torch.cat((input,xy_term.view([-1,1])),1)
        input_new = input_new.to(device,non_blocking=True)   
        
        production_dt = production_weight.mm(torch.transpose(input_new,0,1))
        decay_dt = decay_weight.mm(torch.transpose(input_new,0,1))
        cross_talk_dt = cross_talk_weight.mm(torch.transpose(input_new,0,1))
        hidden_state = production_dt - decay_dt - cross_talk_dt
        out = hidden_state.view([-1,3]) + input
        
        if bias is not None:
            out += bias.unsqueeze(0).expand_as(out)
            
        ctx.save_for_backward(input_new,production_weight,decay_weight,
                              cross_talk_weight,bias,prod_mask,decay_mask,
                              cross_mask)
        
        return out
    
    
    @staticmethod
    def backward(ctx, grad_out):
                
        input_new,production_weight,decay_weight,cross_talk_weight,bias,prod_mask,decay_mask,cross_mask = ctx.saved_tensors
        grad_input = grad_prod_wt = grad_decay_wt = grad_cross_wt = None
        grad_bias = grad_prod_mask = grad_decay_mask = grad_cross_mask = None        
        
        grad_out_new = torch.cat((grad_out,(grad_out[0][0] + grad_out[0][2]).view([-1,1])),1)
        
        if ctx.needs_input_grad[0]:            
            grad_input = (production_weight - decay_weight - cross_talk_weight).mm(grad_out_new)
        
        if ctx.needs_input_grad[1]:
            grad_prod_wt = grad_out.t().clone().mm(input_new)            
            grad_prod_wt = torch.mul(grad_prod_wt, prod_mask)

            grad_decay_wt = grad_out.t().clone().mm(input_new)
            grad_decay_wt = torch.mul(grad_decay_wt, decay_mask)

            grad_cross_wt = grad_out.t().clone().mm(input_new)
            grad_cross_wt = torch.mul(grad_cross_wt, cross_mask)
                        
#         if ctx.needs_input_grad[2]: #don't train for bias
#             grad_bias = (grad_prod_out + grad_decay_out + grad_cross_out).sum(0).squeeze(0)
            
        return grad_input, grad_prod_wt, grad_decay_wt, grad_cross_wt, grad_bias, grad_prod_mask, grad_decay_mask, grad_cross_mask
        
        
class NeuralNet(nn.Module):
    
    def __init__(self,param, prod_mask,decay_mask,cross_mask, bias=True):
        super(NeuralNet, self).__init__() 
        
        if isinstance(prod_mask, torch.Tensor):
            self.prod_mask = prod_mask.type(torch.float)
        else:
            self.prod_mask = torch.tensor(self.prod_mask,dtype=torch.float)
            
        self.prod_mask = nn.Parameter(self.prod_mask.to(device),requires_grad=False)
        
        if isinstance(decay_mask, torch.Tensor):
            self.decay_mask = decay_mask.type(torch.float)
        else:
            self.decay_mask = torch.tensor(self.decay_mask,dtype=torch.float)
            
        self.decay_mask = nn.Parameter(self.decay_mask.to(device),requires_grad=False)
        
        if isinstance(cross_mask, torch.Tensor):
            self.cross_mask = cross_mask.type(torch.float)
        else:
            self.cross_mask = torch.tensor(self.cross_mask,dtype=torch.float)
            
        self.cross_mask = nn.Parameter(self.cross_mask.to(device),requires_grad=False)
        
        if bias:
            self.bias = Variable(torch.zeros(self.prod_mask.shape[0])).type(dtype)
            self.bias = nn.Parameter(self.bias.to(device), requires_grad=False)
        else:
            self.register_parameter('bias',None)
        
        # param_order = [gamma,alpha_xy,beta_y,alpha0,alpha_y,alpha1,alpha2,alpha3]
        # Production Matrix
        self.w_production = torch.tensor([[param[0], 0.0, 0.0, 0.0],
                                          [param[2], 0.0, 0.0, 0.0],
                                          [0.0, param[3], 0.0, 0.0]])
        self.w_production = Variable(self.w_production).type(dtype)
        self.w_production = nn.Parameter(self.w_production.to(device),
                                         requires_grad=True)
        
        # Degradation Matrix
        self.w_decay = torch.tensor([[param[5], 0.0, 0.0, 0.0],
                                     [0.0, param[6], 0.0, 0.0],
                                     [0.0, 0.0, param[4], 0.0]])
        self.w_decay = Variable(self.w_decay).type(dtype)
        self.w_decay = nn.Parameter(self.w_decay.to(device),
                                    requires_grad = True)        
        
        # Cross-Talk Matrix
        self.w_cross_talk = torch.tensor([[0.0, 0.0, 0.0, param[1]],
                                          [0.0, 0.0, 0.0, 0.0],
                                          [0.0, 0.0, 0.0, param[7]]])
        self.w_cross_talk = Variable(self.w_cross_talk).type(dtype)
        self.w_cross_talk = nn.Parameter(self.w_cross_talk.to(device),
                                         requires_grad = True) 
               
         
        self.w_production.data = torch.mul(self.w_production.data, self.prod_mask)
        self.w_decay.data = torch.mul(self.w_decay.data, self.decay_mask)
        self.w_cross_talk.data = torch.mul(self.w_cross_talk.data, self.cross_mask)
        
        
        
    def forward(self,input): 
        
        return CustomizedForwardBackwardFunc.apply(input,self.w_production,
                                               self.w_decay,self.w_cross_talk,
                                               self.bias,self.prod_mask,
                                               self.decay_mask,
                                               self.cross_mask)
    
    

# params initialized  from lmfit
# init_params = [gamma,alpha_xy,beta_y,alpha0,alpha_y,alpha1,alpha2,alpha3]
# init_params = [2.2, 3.1, 1.1, 1.0, 0.79, 0.07, 0.5, 0.008]
init_params = [2.0,3.7,1.5,1.1,0.9,0.1,0.9,0.01]
prod_mask = torch.tensor([[1.0, 0.0, 0.0, 0.0],
                          [1.0, 0.0, 0.0, 0.0],
                          [0.0, 1.0, 0.0, 0.0]])

decay_mask = torch.tensor([[1.0, 0.0, 0.0, 0.0],
                           [0.0, 1.0, 0.0, 0.0],
                           [0.0, 0.0, 1.0, 0.0]])

cross_mask = torch.tensor([[0.0, 0.0, 0.0, 1.0],
                           [0.0, 0.0, 0.0, 0.0],
                           [0.0, 0.0, 0.0, 1.0]])

# define model in GPU
model = NeuralNet(init_params,prod_mask,decay_mask,cross_mask)
model = model.to(device)

loss_function = nn.MSELoss()

optimizer = torch.optim.SGD(model.parameters(), lr=1e-3,momentum = 0.9)

start = datetime.now()
learning_rate = 1e-6
loss_1 = []

for i in tqdm(range(epochs)):
    total_loss = 0
    
    for j in range(x_train.size(0)):
        optimizer.zero_grad() 
        
#         for param in model.parameters():
#             if param.requires_grad == True:
#                 param.retain_grad()
#                 param.grad.data.zero_() #commented because of error of  
                                          #getting "None" from param.grad
        
        x1 = x_train[j:(j+1)]
        y1 = y_train[j:(j+1)]         
        x1 = x1.to(device,non_blocking=True)
        y1 = y1.to(device,non_blocking=True)                
        pred = model(x1)
               
        loss = loss_function(pred,y1)
        total_loss += loss               
        loss.backward()            
#         optimizer.step() # commented because we're using customized 
                           # gradient upgradation
        with torch.no_grad():
            for name, param in model.named_parameters():                
                if param.requires_grad == True:
                    param -= learning_rate * param.grad
                                        
                    '''
                        Print statement to cross-cross check the masked
                        gradradients
                    '''
#                     print(param.grad)
                    # check masked param.grad
#                     if np.array(param.grad.cpu()).size == np.array(prod_mask).size:
#                         print('--- epoch={}, loss={} ---'.format(t,loss.item()))
#                         print('↓↓↓masked weight↓↓↓')
#                         print(param)
#                         print('↓↓↓masked grad of weight↓↓↓')
#                         print(param.grad)
             
    
    # Early Stopping
    loss_1.append(total_loss.item())  
#     print(total_loss.item())
    if len(loss_1) &gt; 1 and np.abs(loss_1[-2] - total_loss.item()) &lt;= 0.01:
        print('Early Stopping due to dIfference in error is less than threshold (0.10)')
        print('Stopping training at %d epochs with total loss = %d' %(i,total_loss))
        break


    if i % 50 == 0:
        print("Epoch: {} loss {}".format(i, total_loss))
        
end = datetime.now()
time_taken = end - start
print('Execution Time: ',time_taken)
</code></pre>
<p>First I tried with <code>learning_rate = 1e-3</code> loss was continuously increasing. Then I tried <code>1e-4</code> still the same model response. Then I tried <code>1e-5</code> then model loss decreases to 14 epochs but then started increasing.  So I further lowered down the learning rate to <code>1e-6</code> then the model was performing well. I also used early stopping criteria of difference in error as 0.1</p>
<p>Here is the model prediction results:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/1/f1f908f98bac7298d89535335671f1c65e445d6c.png" data-download-href="https://discuss.pytorch.org/uploads/default/f1f908f98bac7298d89535335671f1c65e445d6c" title="result1_original_param"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/1/f1f908f98bac7298d89535335671f1c65e445d6c_2_690x198.png" alt="result1_original_param" data-base62-sha1="ywACMx0ZERHumNPC7YQrUXJDDxq" width="690" height="198" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/1/f1f908f98bac7298d89535335671f1c65e445d6c_2_690x198.png, https://discuss.pytorch.org/uploads/default/original/3X/f/1/f1f908f98bac7298d89535335671f1c65e445d6c.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/f/1/f1f908f98bac7298d89535335671f1c65e445d6c.png 2x" data-dominant-color="F7F6F5"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">result1_original_param</span><span class="informations">877×252 33.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>My issues are as follows:</p>
<ol>
<li>With learning rate 1e-6 and early stopping criteria (i.e. the difference in epoch error less than 0.1) then error lowered up to 1880 and then model training stopped. (Ideally, it must go down to zero). Why??</li>
<li>I lowered the early stopping condition to 0.01 with the same learning rate 1e-6 then the model performed well till 150 epochs (error at 150 epoch = 1880) and then error started increasing again.</li>
</ol>
<p>I though lowering the error criteria will help the model but opposite is happening. My doubt is :</p>
<ol>
<li>Is there anything wrong in the model that error is not going below 1880.</li>
<li>I am taking just one batch of all training samples. How does batch_size affect model training perofrmance.</li>
</ol>

username: VIVEK_RUHELA
cooked: <p>@ <a href="https://discuss.pytorch.org/u/uchida-takumi">uchida-takumi</a> @ <a href="https://discuss.pytorch.org/u/jamesproud">jamesproud</a> @ <a href="https://discuss.pytorch.org/u/Arvind_Subramaniam">Arvind_Subramaniam</a> : What will you suggest. Thanks</p>