优化器参数分配采用 itertools.chain()
多损失回传除最后一个loss.backward(),其余需要添加 retain_graph=True
测试:
import torch
import itertools
import torch.nn as nn
import torch.nn.functional as F
class test_model(nn.Module):
def __init__(self):
super(test_model, self).__init__()
self.fc1 = nn.Linear(5, 4)
self.fc2 = nn.Linear(4, 3)
self.fc3 = nn.Linear(4, 3)
def forward(self, x):
x = self.fc1(x)
out1 = self.fc2(x)
out2 = self.fc3(x)
return out1, out2
x = torch.randn((3, 5))
y = torch.torch.randint(3, (3,), dtype=torch.int64)
model = test_model()
model.train()
optim1 = torch.optim.RMSprop(itertools.chain(model.fc1.parameters(), model.fc2.parameters()), lr=0.001)
optim2 = torch.optim.RMSprop(model.fc3.parameters(), lr=0.001)
print('-----------------------------------------------------------------------')
print("更新前")
print('-----------------------------------------------------------------------')
print(model.fc1.weight)
print(model.fc2.weight)
print(model.fc3.weight)
for i in range(5):
out1, out2 = model(x)
loss1 = F.cross_entropy(out1, y)
loss2 = F.cross_entropy(out2, y)
optim1.zero_grad()
optim2.zero_grad()
loss1.backward(retain_graph=True)
loss2.backward()
optim1.step()
optim2.step()
print('-----------------------------------------------------------------------')
print("更新后")
print('-----------------------------------------------------------------------')
print(model.fc1.weight)
print(model.fc2.weight)
print(model.fc3.weight)
输出:
若loss1.backward()不添加retain_graph=True则会报错,原因是Pytorch每次调用.backward()都会free掉所有buffers,retain_graph=True可以将前一次的backward()的梯度保存在buffer内。