Pytorch required_grad=False does not freeze network parameters when running on GPU

Question

I'm trying to freeze a layer of a toy model when training using Pytorch. In the following code, when I run the code on CPU, the layer isn't updated. (Please see the code line print("%.8f" % np.max(np.abs(before -after)))). However, when I run the code on GPU, the layer is updated. What is wrong with my implementation?

import torch
from torch import nn
from torch.autograd import Variable
import torch.nn.functional as F
import torch.optim as optim
import numpy as np

def toNP(x):
    return x.detach().to('cpu').numpy()

# toy feed-forward net
class Sub_Net(nn.Module):
    def __init__(self):
        super(Sub_Net, self).__init__()

        self.fc1 = nn.Linear(10, 3)
        self.fc2 = nn.Linear(3, 3)
        self.fc3 = nn.Linear(3, 3)
        self.fc4 = nn.Linear(3, 3)
        self.fc5 = nn.Linear(3, 1)

    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        x = self.fc4(x)
        x = self.fc5(x)
        return x

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()

        self.Sub_Net = Sub_Net()

        self.fc1 = nn.Linear(10, 3)
        self.fc2 = nn.Linear(3, 3)
        self.fc3 = nn.Linear(3, 3)
        self.fc4 = nn.Linear(3, 3)
        self.fc5 = nn.Linear(3, 1)

    def forward(self, x):

        y = self.Sub_Net(x)

        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        x = self.fc4(x)
        x = self.fc5(x)
        return x+y

def generator_step(net, optimizer, criterion, input, target):

    output = net(input)
    loss = criterion(output, target)
    net.zero_grad()
    loss.backward()
    optimizer.step()

def discrimination_step(net, optimizer, criterion, input, target):

    for param in net.Sub_Net.parameters():
        param.requires_grad = False

    before = toNP(net.Sub_Net.fc2.weight)

    output = net(input)
    loss = criterion(output, target)
    net.zero_grad()
    loss.backward()
    optimizer.step()

    after = toNP(net.Sub_Net.fc2.weight)

    print("%.8f" % np.max(np.abs(before -after)) )


# Run model on GPU
# net = Net().type(torch.cuda.FloatTensor)
# random_input = Variable(torch.randn(10, )).cuda()
# random_target = Variable(torch.randn(1, )).cuda()

# Run model on CPU
net = Net()
random_input = Variable(torch.randn(10, ))
random_target = Variable(torch.randn(1, ))

# loss
criterion = nn.MSELoss()
optimizer = optim.Adam(filter(lambda p: p.requires_grad, net.parameters()), lr=0.1)

for epoch in range(1, 10):

    generator_step(net, optimizer, criterion, random_input, random_target)
    discrimination_step(net, optimizer, criterion, random_input, random_target)

Results when running on CPU

Results when running on GPU

user_17189149 · Answer

It seems it's not CPU / GPU acting differently, but it is about .to('cpu') function in your toNP function.

If given tensor is on GPU, it returns copied tensor on cpu, while it returns given original object when the given tensor is already on CPU. Please refer more on this site.

To clarify, I've added print function to your discriminator_step as follows:

def discrimination_step(net, optimizer, criterion, input, target):

    for param in net.Sub_Net.parameters():
        param.requires_grad = False

    before = toNP(net.Sub_Net.fc2.weight)
    print(f'Before:
{before}')

    output = net(input)
    loss = criterion(output, target)
    net.zero_grad()
    loss.backward()
    optimizer.step()

    print(f'Before:
{before}')
    after = toNP(net.Sub_Net.fc2.weight)
    print(f'After:
{after}')

    print("dff: %.8f" % np.max(np.abs(before -after)) )

Then the code yields results when on CPU (for 1 epoch):

Before:
[[-0.0222426   0.06449176  0.41833472]
 [-0.3276776  -0.22486973  0.38021228]
 [-0.37726757  0.26268137 -0.05000275]]
Before:
[[ 0.04476321  0.13149747  0.48534054]
 [-0.2606718  -0.15786391  0.44721812]
 [-0.31026173  0.32968715  0.01700307]]
After:
[[ 0.04476321  0.13149747  0.48534054]
 [-0.2606718  -0.15786391  0.44721812]
 [-0.31026173  0.32968715  0.01700307]]
dff: 0.00000000

and on GPU:

Before:
[[-0.06808002  0.39740798  0.55723506]
 [-0.17421165 -0.36702433 -0.4208245 ]
 [-0.37865937 -0.52346057 -0.15856335]]
Before:
[[-0.06808002  0.39740798  0.55723506]
 [-0.17421165 -0.36702433 -0.4208245 ]
 [-0.37865937 -0.52346057 -0.15856335]]
After:
[[-0.13508584  0.4644138   0.6242409 ]
 [-0.24121748 -0.30001852 -0.35381868]
 [-0.31165355 -0.5904664  -0.22556916]]
dff: 0.06700583

It shows before values are changed since the returned before tensor shares the same storage with net.Sub_Net.fc2.weight when on CPU. Layers are updated regardless of CPU and GPU since they are already in parameter groups of the Adam optimizer.

user_10475762 · Answer

(See edit for the actual answer to the question, the top question is still relevant to those experiencing this problem in a broader sense)

Seems directly related to this link.

If you take the param grad removal out of the forward call, it works:

def discrimination_step(net, optimizer, criterion, input, target):     
    before = toNP(net.Sub_Net.fc2.weight)      
    output = net(input) 
    loss = criterion(output, target) 
    net.zero_grad() 
    loss.backward() 
    optimizer.step() 
 
    after = toNP(net.Sub_Net.fc2.weight)      
    print("%.8f" % np.max(np.abs(before -after)) ) 
 
# Run model on GPU 
net = Net().type(torch.cuda.FloatTensor) 
random_input = Variable(torch.randn(10, )).cuda() 
random_target = Variable(torch.randn(1, )).cuda() 
for param in net.Sub_Net.parameters():  
    param.requires_grad = False  
# Run model on CPU 
#net = Net() 
#random_input = Variable(torch.randn(10, )) 
#random_target = Variable(torch.randn(1, )) 
...
 
for epoch in range(1, 10): 
 
    generator_step(net, optimizer, criterion, random_input, random_target) 
    discrimination_step(net, optimizer, criterion, random_input, random_target)

Which gives the output:

EDIT

You can just set the grad to None during the forward pass, which makes the Adam optimizer ignore it. If you just use zero grad Adam has moving averages which still change the values (I also checked this as SGD works with just net.Sub_Net.zero_grad()).

def generator_step(net, optimizer, criterion, input, target):
    beforeG = toNP(net.fc2.weight)
    beforeD = toNP(net.Sub_Net.fc2.weight)
    output = net(input)
    loss = criterion(output, target)
    net.zero_grad()
    loss.backward()
    optimizer.step()
    afterG = toNP(net.fc2.weight)
    afterD = toNP(net.Sub_Net.fc2.weight)
    print("gen step - G %.8f" % np.max(np.abs(beforeG -afterG)) )
    print("gen step - D %.8f" % np.max(np.abs(beforeD -afterD)) )

def discrimination_step(net, optimizer, criterion, input, target):
    beforeG = toNP(net.fc2.weight)
    beforeD = toNP(net.Sub_Net.fc2.weight)
    output = net(input)
    loss = criterion(output, target)
    net.zero_grad()
    loss.backward()
    for param in net.Sub_Net.parameters():
        param.grad = None
    #print(net.fc3.weight.grad)
    #print(net.Sub_Net.fc3.weight.grad)
    optimizer.step()
    afterG = toNP(net.fc2.weight)
    afterD = toNP(net.Sub_Net.fc2.weight)
    print("disc step - G %.8f" % np.max(np.abs(beforeG -afterG)) )
    print("disc step - D %.8f" % np.max(np.abs(beforeD -afterD)) )

# Run model on GPU
net = Net().type(torch.cuda.FloatTensor)
random_input = Variable(torch.randn(10, )).cuda()
random_target = Variable(torch.randn(1, )).cuda()

criterion = nn.MSELoss()
optimizer = optim.SGD([i for i in net.parameters()], lr=0.1)
for epoch in range(1, 3):
    generator_step(net, optimizer, criterion, random_input, random_target)
    discrimination_step(net, optimizer, criterion, random_input, random_target)

optimizer = optim.Adam([i for i in net.parameters()], lr=0.1)
for epoch in range(1, 3):
    generator_step(net, optimizer, criterion, random_input, random_target)
    discrimination_step(net, optimizer, criterion, random_input, random_target)

Outputs:

gen step - G 0.02727062
gen step - D 0.03749144
disc step - G 0.00164142
disc step - D 0.00000000
gen step - G 0.00133285
gen step - D 0.00459737
disc step - G 0.00046796
disc step - D 0.00000000

gen step - G 0.09999971
gen step - D 0.09999990
disc step - G 0.08018596
disc step - D 0.00000000
gen step - G 0.08054254
gen step - D 0.07518816
disc step - G 0.06967809
disc step - D 0.00000000

Pytorch required_grad=False does not freeze network parameters when running on GPU

2 Answers