Khah_Nguyen:

ew in pytorch and I have met some problems in updating parameters


@Khah_Nguyen you will have to manually turn on and off requires_grad while doing the backward. This will ensure that when the optimizer step function runs, the appropriate grads are stored in the node
import numpy as np
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable

class random_model(nn.Module):
    def __init__(self, num_layers):
        super(random_model, self).__init__()
        self.model1 = nn.Sequential(nn.Linear(100, 20), nn.BatchNorm1d(20), nn.ReLU())
        self.model3 = nn.Linear(20, 1)
    
    def forward(self, x):
        x = self.model1(x)
        x = self.model3(x)
        return x
    
model = random_model(10)
loss1 = torch.nn.MSELoss()
loss2 = torch.nn.SmoothL1Loss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

X = torch.rand(100, 100)
y = torch.rand(100)

for cur_epoch in range(100):
    optimizer.zero_grad()
    output = model(X)
    cur_loss1 = loss1(output, y)
    cur_loss2 = loss2(output, y)
    cur_loss = cur_loss1 + cur_loss2
    for param in model.model3.parameters():
        param.requires_grad = False
    
    cur_loss.backward(retain_graph=True)
    for param in model.model1.parameters():
        param.requires_grad = False
    for param in model.model3.parameters():
        param.requires_grad = True
    cur_loss2.backward()
    optimizer.step()
    print("Epoch {0} Loss is {1}".format(cur_epoch, cur_loss.item()))