Pytorch 同一个网络同时拥有多个输出的写法

文章目录 展开

Pytorch 是非常热门的神经网络计算框架，一般来说一个网络对应一个输出，如果我们需要一个网络同时输出多个 output，这在 Pytorch 也是可以做到的。输出之后，可以分别计算 Loss，然后进行 BP。

一、同一个网络输出多个结果

只需要在 forward() 部分这么写：

def forward(self, x):
    # Do your stuff here
    ...
    x1 = F.log_softmax(x) # class probabilities
    x2 = ... # bounding box calculation
    return x1, x2

然后分别计算 Loss 并训练网络：

out1, out2 = model(data)
loss1 = criterion1(out1, target1)
loss2 = criterion2(out2, target2)
loss = loss1 + loss2
loss.backward()

具体的话就根据自己的具体需求来改了。

参考：https://discuss.pytorch.org/t/a-model-with-multiple-outputs/10440/2

二、其他参考

这里还有一份另外一个同学的提问，不过暂时没人答复：https://discuss.pytorch.org/t/training-a-multiple-output-network/6968

看了下总体思想是一样的，应该没什么问题。

import torch
import torchvision
import torch.nn as nn
import numpy as np
import torch.utils.data as data
import torchvision.transforms as transforms
import torchvision.datasets as dsets
from torch.autograd import Variable
import pdb


class RegNet(nn.Module):
  def __init__(self):
    super(RegNet,self).__init__()
    self.Reg_x = nn.Linear(4,4)
    self.Reg_y = nn.Linear(4,3)
    
  def forward(self,x):
    x_x = self.Reg_x(x)
    x_y = self.Reg_y(x)
    return x_x,x_y


# Input
z = Variable(torch.ones(5,4),requires_grad=True)

# Rep Layer
Rep = nn.Linear(4,4)

# Reg Layer
Reg = RegNet()

label_x = Variable(torch.ones(5,4))
label_y = Variable(torch.zeros(5,3))

criterion = nn.MSELoss()
optimizer1 = torch.optim.SGD(Rep.parameters(), lr=0.01)
optimizer2 = torch.optim.SGD(Reg.parameters(), lr=0.01)

# train
# feed-forward
rep_out = Rep(z)
pred_x,pred_y = Reg(rep_out)

loss_x = criterion(pred_x,label_x)
loss_y = criterion(pred_y,label_y)
loss = loss_x + loss_y

loss.backward()