?在pytorch中停止梯度流的若干辦法，避免不必要模塊的參數(shù)更新

2020/4/11 FesianXu

前言

在現(xiàn)在的深度模型軟件框架中，如TensorFlow和PyTorch等等，都是實現(xiàn)了自動求導(dǎo)機制的。在深度學(xué)習(xí)中，有時候我們需要對某些模塊的梯度流進(jìn)行精確地控制，包括是否允許某個模塊的參數(shù)更新，更新地幅度多少，是否每個模塊更新地幅度都是一樣的。這些問題非常常見，但是在實踐中卻很容易出錯，我們在這篇文章中嘗試對第一個子問題，也就是如果精確控制某些模型是否允許其參數(shù)更新，進(jìn)行總結(jié)。如有謬誤，請聯(lián)系指出，轉(zhuǎn)載請注明出處。

本文實驗平臺：pytorch 1.4.0, ubuntu 18.04, python 3.6

聯(lián)系方式：

e-mail: FesianXu@gmail.com

QQ: 973926198

github: https://github.com/FesianXu

知乎專欄: 計算機視覺/計算機圖形理論與應(yīng)用

微信公眾號：

---

為什么我們要控制梯度流

為什么我們要控制梯度流？這個答案有很多個，但是都可以歸結(jié)為避免不需要更新的模型模塊被參數(shù)更新。我們在深度模型訓(xùn)練過程中，很可能存在多個loss，比如GAN對抗生成網(wǎng)絡(luò)，存在G_loss和D_loss，通常來說，我們通過D_loss只希望更新判別器(Discriminator)，而生成網(wǎng)絡(luò)(Generator)并不需要，也不能被更新；生成網(wǎng)絡(luò)只在通過G_loss學(xué)習(xí)的情況下，才能被更新。這個時候，如果我們不控制梯度流，那么我們在訓(xùn)練D_loss的時候，我們的前端網(wǎng)絡(luò)Generator和CNN難免也會被一起訓(xùn)練，這個是我們不期望發(fā)生的。

Fig 1.1 典型的GAN結(jié)構(gòu)，由生成器和判別器組成。

多個loss的協(xié)調(diào)只是其中一種情況，還有一種情況是：我們在進(jìn)行模型遷移的過程中，經(jīng)常采用某些已經(jīng)預(yù)訓(xùn)練好了的特征提取網(wǎng)絡(luò)，比如VGG, ResNet之類的，在適用到具體的業(yè)務(wù)數(shù)據(jù)集時候，特別是小數(shù)據(jù)集的時候，我們可能會希望這些前端的特征提取器不要更新，而只是更新末端的分類器(因為數(shù)據(jù)集很小的情況下，如果貿(mào)然更新特征提取器，很可能出現(xiàn)不期望的嚴(yán)重過擬合，這個時候的合適做法應(yīng)該是更新分類器優(yōu)先)，這個時候我們也可以考慮停止特征提取器的梯度流。

這些情況還有很多，我們在實踐中發(fā)現(xiàn)，精確控制某些模塊的梯度流是非常重要的。筆者在本文中打算討論的是對某些模塊的梯度流的截斷，而并沒有討論對某些模塊梯度流的比例縮放，或者說最細(xì)粒度的梯度流控制，后者我們將會在后文中討論。

一般來說，截斷梯度流可以有幾種思路：

停止計算某個模塊的梯度，在優(yōu)化過程中這個模塊還是會被考慮更新，然而因為梯度已經(jīng)被截斷了，因此不能被更新。

• 設(shè)置tensor.detach()：完全截斷之前的梯度流
• 設(shè)置參數(shù)的requires_grad屬性：單純不計算當(dāng)前設(shè)置參數(shù)的梯度，不影響梯度流
• torch.no_grad()：效果類似于設(shè)置參數(shù)的requires_grad屬性

在優(yōu)化器中設(shè)置不更新某個模塊的參數(shù)，這個模塊的參數(shù)在優(yōu)化過程中就不會得到更新，然而這個模塊的梯度在反向傳播時仍然可能被計算。

我們后面分別按照這兩大類思路進(jìn)行討論。

---

停止計算某個模塊的梯度

在本大類方法中，主要涉及到了tensor.detach()和requires_grad的設(shè)置，這兩種都無非是對某些模塊，某些節(jié)點變量設(shè)置了是否需要梯度的選項。

tensor.detach()

tensor.detach()的作用是：

tensor.detach()會創(chuàng)建一個與原來張量共享內(nèi)存空間的一個新的張量，不同的是，這個新的張量將不會有梯度流流過，這個新的張量就像是從原先的計算圖中脫離(detach)出來一樣，對這個新的張量進(jìn)行的任何操作都不會影響到原先的計算圖了。因此對此新的張量進(jìn)行的梯度流也不會流過原先的計算圖，從而起到了截斷的目的。

這樣說可能不夠清楚，我們舉個例子。眾所周知，我們的pytorch是動態(tài)計算圖網(wǎng)絡(luò)，正是因為計算圖的存在，才能實現(xiàn)自動求導(dǎo)機制。考慮一個表達(dá)式：

如果用計算圖表示則如Fig 2.1所示。

Fig 2.1 計算圖示例

考慮在這個式子的基礎(chǔ)上，加上一個分支：

那么計算圖就變成了：

Fig 2.2 添加了新的分支后的計算圖

如果我們不detach() 中間的變量z，分別對pq和w進(jìn)行反向傳播梯度，我們會有：

x?=?torch.tensor(([1.0]),requires_grad=True)
y?=?x**2
z?=?2*y
w=?z**3

#?This?is?the?subpath
#?Do?not?use?detach()
p?=?z
q?=?torch.tensor(([2.0]),?requires_grad=True)
pq?=?p*q
pq.backward(retain_graph=True)

w.backward()
print(x.grad)

輸出結(jié)果為 tensor([56.])。我們發(fā)現(xiàn)，這個結(jié)果是吧pq和w的反向傳播結(jié)果都進(jìn)行了考慮的，也就是新增加的分支的反向傳播影響了原先主要枝干的梯度流。這個時候我們用detach()可以把p給從原先計算圖中脫離出來，使得其不會干擾原先的計算圖的梯度流，如：

Fig 2.3 用了detach之后的計算圖

那么，代碼就對應(yīng)地修改為：

x?=?torch.tensor(([1.0]),requires_grad=True)
y?=?x**2
z?=?2*y
w=?z**3

#?detach?it,?so?the?gradient?w.r.t?`p`?does?not?effect?`z`!
p?=?z.detach()
q?=?torch.tensor(([2.0]),?requires_grad=True)
pq?=?p*q
pq.backward(retain_graph=True)
w.backward()
print(x.grad)

這個時候，因為分支的梯度流已經(jīng)影響不到原先的計算圖梯度流了，因此輸出為tensor([48.])。

這只是個計算圖的簡單例子，在實際模塊中，我們同樣可以這樣用，舉個GAN的例子，代碼如：

????def?backward_D(self):
????????#?Fake
????????#?stop?backprop?to?the?generator?by?detaching?fake_B
????????fake_AB?=?self.fake_B
????????#?fake_AB?=?self.fake_AB_pool.query(torch.cat((self.real_A,?self.fake_B),?1))
????????self.pred_fake?=?self.netD.forward(fake_AB.detach())
????????self.loss_D_fake?=?self.criterionGAN(self.pred_fake,?False)

????????#?Real
????????real_AB?=?self.real_B?#?GroundTruth
????????#?real_AB?=?torch.cat((self.real_A,?self.real_B),?1)
????????self.pred_real?=?self.netD.forward(real_AB)
????????self.loss_D_real?=?self.criterionGAN(self.pred_real,?True)

????????#?Combined?loss
????????self.loss_D?=?(self.loss_D_fake?+?self.loss_D_real)?*?0.5

????????self.loss_D.backward()

????def?backward_G(self):
????????#?First,?G(A)?should?fake?the?discriminator
????????fake_AB?=?self.fake_B
????????pred_fake?=?self.netD.forward(fake_AB)
????????self.loss_G_GAN?=?self.criterionGAN(pred_fake,?True)

????????#?Second,?G(A)?=?B
????????self.loss_G_L1?=?self.criterionL1(self.fake_B,?self.real_B)?*?self.opt.lambda_A

????????self.loss_G?=?self.loss_G_GAN?+?self.loss_G_L1

????????self.loss_G.backward()


????def?forward(self):
????????self.real_A?=?Variable(self.input_A)
????????self.fake_B?=?self.netG.forward(self.real_A)
????????self.real_B?=?Variable(self.input_B)

????#?先調(diào)用 forward, 再 D backward，?更新D之后；?再G backward，?再更新G
????def?optimize_parameters(self):
????????self.forward()

????????self.optimizer_D.zero_grad()
????????self.backward_D()
????????self.optimizer_D.step()

????????self.optimizer_G.zero_grad()
????????self.backward_G()
????????self.optimizer_G.step()

我們注意看第六行，self.pred_fake = self.netD.forward(fake_AB.detach())使得在反向傳播D_loss的時候不會更新到self.netG，因為fake_AB是由self.netG生成的，代碼如self.fake_B = self.netG.forward(self.real_A)。

設(shè)置requires_grad

tensor.detach()是截斷梯度流的一個好辦法，但是在設(shè)置了detach()的張量之前的所有模塊，梯度流都不能回流了(不包括這個張量本身，這個張量已經(jīng)脫離原先的計算圖了)，如以下代碼所示：

x?=?torch.randn(2,?2)
x.requires_grad?=?True

lin0?=?nn.Linear(2,?2)
lin1?=?nn.Linear(2,?2)
lin2?=?nn.Linear(2,?2)
lin3?=?nn.Linear(2,?2)
x1?=?lin0(x)
x2?=?lin1(x1)
x2?=?x2.detach()?#?此處設(shè)置了detach，之前的所有梯度流都不會回傳了
x3?=?lin2(x2)
x4?=?lin3(x3)
x4.sum().backward()
print(lin0.weight.grad)
print(lin1.weight.grad)
print(lin2.weight.grad)
print(lin3.weight.grad)

輸出為:

None
None
tensor([[-0.7784,?-0.7018],
????????[-0.4261,?-0.3842]])
tensor([[?0.5509,?-0.0386],
????????[?0.5509,?-0.0386]])

我們發(fā)現(xiàn)lin0.weight.grad和lin0.weight.grad都為None了，因為通過脫離中間張量，原先計算圖已經(jīng)和當(dāng)前回傳的梯度流脫離關(guān)系了。

這樣有時候不夠理想，因為我們可能存在只需要某些中間模塊不計算梯度，但是梯度仍然需要回傳的情況，在這種情況下，如下圖所示，我們可能只需要不計算B_net的梯度，但是我們又希望計算A_net和C_net的梯度，這個時候怎么辦呢？當(dāng)然，通過detach()這個方法是不能用了。

事實上，我們可以通過設(shè)置張量的requires_grad屬性來設(shè)置某個張量是否計算梯度，而這個不會影響梯度回傳，只會影響當(dāng)前的張量。修改上面的代碼，我們有：

x?=?torch.randn(2,?2)
x.requires_grad?=?True

lin0?=?nn.Linear(2,?2)
lin1?=?nn.Linear(2,?2)
lin2?=?nn.Linear(2,?2)
lin3?=?nn.Linear(2,?2)
x1?=?lin0(x)
x2?=?lin1(x1)
for?p?in?lin2.parameters():
????p.requires_grad?=?False
x3?=?lin2(x2)
x4?=?lin3(x3)
x4.sum().backward()
print(lin0.weight.grad)
print(lin1.weight.grad)
print(lin2.weight.grad)
print(lin3.weight.grad)

輸出為:

tensor([[-0.0117,??0.9976],
????????[-0.0080,??0.6855]])
tensor([[-0.0075,?-0.0521],
????????[-0.0391,?-0.2708]])
None
tensor([[0.0523,?0.5429],
????????[0.0523,?0.5429]])

啊哈，正是我們想要的結(jié)果，只有設(shè)置了requires_grad=False的模塊沒有計算梯度，但是梯度流又能夠回傳。

另外，設(shè)置requires_grad經(jīng)常用在對輸入變量和輸入的標(biāo)簽進(jìn)行新建的時候使用，如：

for?mat,label?in?dataloader:
????mat?=?Variable(mat,?requires_grad=False)
????label?=?Variable(mat,requires_grad=False)
????...

當(dāng)然，通過把所有前端網(wǎng)絡(luò)都設(shè)置requires_grad=False，我們可以實現(xiàn)類似于detach()的效果，也就是把該節(jié)點之前的所有梯度流回傳截斷。以VGG16為例子，如果我們只需要訓(xùn)練其分類器，而固定住其特征提取器網(wǎng)絡(luò)的參數(shù)，我們可以采用將前端網(wǎng)絡(luò)的所有參數(shù)的requires_grad設(shè)置為False，因為這個時候完全不需要梯度流的回傳，只需要前向計算即可。代碼如：

model?=?torchvision.models.vgg16(pretrained=True)
for?param?in?model.features.parameters():
????param.requires_grad?=?False

torch.no_grad()

在對訓(xùn)練好的模型進(jìn)行評估測試時，我們同樣不需要訓(xùn)練，自然也不需要梯度流信息了。我們可以把所有參數(shù)的requires_grad屬性設(shè)置為False，事實上，我們常用torch.no_grad()上下文管理器達(dá)到這個目的。即便輸入的張量屬性是requires_grad=True, ? torch.no_grad()可以將所有的中間計算結(jié)果的該屬性臨時轉(zhuǎn)變?yōu)镕alse。

如例子所示：

x?=?torch.randn(3,?requires_grad=True)
x1?=?(x**2)
print(x.requires_grad)
print(x1.requires_grad)

with?torch.no_grad():
????x2?=?(x**2)
????print(x1.requires_grad)
????print(x2.requires_grad)

輸出為：

True
True
True
False

注意到只是在torch.no_grad()上下文管理器范圍內(nèi)計算的中間變量的屬性requires_grad才會被轉(zhuǎn)變?yōu)镕alse，在該管理器外面計算的并不會變化。

不過和單純手動設(shè)置requires_grad=False不同的是，在設(shè)置了torch.no_grad()之前的層是不能回傳梯度的，延續(xù)之前的例子如：

x?=?torch.randn(2,?2)
x.requires_grad?=?True

lin0?=?nn.Linear(2,?2)
lin1?=?nn.Linear(2,?2)
lin2?=?nn.Linear(2,?2)
lin3?=?nn.Linear(2,?2)
x1?=?lin0(x)
with?torch.no_grad():
????x2?=?lin1(x1)
x3?=?lin2(x2)
x4?=?lin3(x3)
x4.sum().backward()
print(lin0.weight.grad)
print(lin1.weight.grad)
print(lin2.weight.grad)
print(lin3.weight.grad)

輸出為：

None
None
tensor([[-0.0926,?-0.0945],
????????[-0.2793,?-0.2851]])
tensor([[-0.5216,??0.8088],
????????[-0.5216,??0.8088]])

此處如果我們打印lin1.weight.requires_grad我們會發(fā)現(xiàn)其為True，但是其中間變量x2.requires_grad=False。

一般來說在實踐中，我們的torch.no_grad()通常會在測試模型的時候使用，而不會選擇在選擇性訓(xùn)練某些模塊時使用[1]，例子如：

model.train()
#?here?train?the?model,?just?skip?the?codes
model.eval()?#?here?we?start?to?evaluate?the?model
with?torch.no_grad():
?for?each?in?eval_data:
??data,?label?=?each
??logit?=?model(data)
??...?#?here?we?just?skip?the?codes

注意

通過設(shè)置屬性requires_grad=False的方法(包括torch.no_grad())很多時候可以避免保存中間計算的buffer，從而減少對內(nèi)存的需求，但是這個也是視情況而定的，比如如[2]的所示

graph LR;
input-->A_net;
A_net-->B_net;
B_net-->C_net;

如果我們不需要A_net的梯度，我們設(shè)置所有A_net的requires_grad=False，因為后續(xù)的B_net和C_net的梯度流并不依賴于A_net，因此不計算A_net的梯度流意味著不需要保存這個中間計算結(jié)果，因此減少了內(nèi)存。

但是如果我們不需要的是B_net的梯度，而需要A_net和C_net的梯度，那么問題就不一樣了，因為A_net梯度依賴于B_net的梯度，就算不計算B_net的梯度，也需要保存回傳過程中B_net中間計算的結(jié)果，因此內(nèi)存并不會被減少。

但是通過tensor.detach()的方法并不會減少內(nèi)存使用，這一點需要注意。

---

設(shè)置優(yōu)化器的更新列表

這個方法更為直接，即便某個模塊進(jìn)行了梯度計算，我只需要在優(yōu)化器中指定不更新該模塊的參數(shù)，那么這個模塊就和沒有計算梯度有著同樣的效果了。如以下代碼所示:

class?model(nn.Module):
????def?__init__(self):
????????super().__init__()
????????self.model_1?=?nn.linear(10,10)
????????self.model_2?=?nn.linear(10,20)
????????self.fc?=?nn.linear(20,2)
????????self.relu?=?nn.ReLU()
???????
????def?foward(inputv):
????????h?=?self.model_1(inputv)
????????h?=?self.relu(h)
????????h?=?self.model_2(inputv)
????????h?=?self.relu(h)
????????return?self.fc(h)

在設(shè)置優(yōu)化器時，我們只需要更新fc層和model_2層，那么則是:

curr_model?=?model()
opt_list?=?list(curr_model.fc.parameters())+list(curr_model.model_2.parameters())
optimizer?=?torch.optim.SGD(opt_list,?lr=1e-4)

當(dāng)然你也可以通過以下的方法去設(shè)置每一個層的學(xué)習(xí)率來避免不需要更新的層的更新[3]：

optim.SGD([
????????????????{'params':?model.model_1.parameters()},
????????????????{'params':?model.mode_2.parameters(),?'lr':?0},
?????????{'params':?model.fc.parameters(),?'lr':?0}
????????????],?lr=1e-2,?momentum=0.9)

這種方法不需要更改模型本身結(jié)構(gòu)，也不需要添加模型的額外節(jié)點，但是需要保存梯度的中間變量，并且將會計算不需要計算的模塊的梯度(即便最后優(yōu)化的時候不考慮更新)，這樣浪費了內(nèi)存和計算時間。

Reference

[1]. https://blog.csdn.net/LoseInVain/article/details/82916163

[2]. https://discuss.pytorch.org/t/requires-grad-false-does-not-save-memory/21936

[3]. https://pytorch.org/docs/stable/optim.html#module-torch.optim

總結(jié)

以上是生活随笔為你收集整理的更新fielddata为true_在pytorch中停止梯度流的若干办法，避免不必要模块的参数更新...的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。