【损失优化】pytorch中多优化器参数分配与多损失回传(解决报错：RuntimeError: Trying to backward through the graph a second time)

优化器参数分配采用 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内。