使用Pytorch训练two-head网络

使用Pytorch训练two-head网络

之前有写过一篇如何使用Pytorch实现two-head（多输出）模型
在那篇文章里，基本把two-head网络以及构建讲清楚了（如果不清楚请先移步至那一篇博文）。但是我后来发现之前的训练方法貌似有些问题。

以前的训练方法：

之前是把两个head分开进行训练的，因此每一轮训练先要对一个batch的数据进行划分，然后再分别训练两个头。代码如下：

			f_out_y0, _ = net(x0)
            _, f_out_y1 = net(x1)

            #实例化损失函数
            criterion0 = Loss()
            criterion1 = Loss()
            loss0 = criterion0(f_y0, f_out_y0, w0)
            loss1 = criterion1(f_y1, f_out_y1, w1)
            print(loss0.item(), loss1.item())
            #对网络参数进行初始化
            optimizer.zero_grad()
            loss0.backward()
            loss1.backward()

            #对网络的参数进行更新
            optimizer.step()

但是在实际操作中想到那这样的话岂不是每次都先使用t=0的数据训练公共的表示层，再使用t=1的数据去训练。这样会不会使表示层产生bias呢？且这样两步训练也很麻烦。

修改后的方法

使用之前训练方法其实还是对神经网络的训练的机理不清楚。事实上，在计算loss的时候每个数据点的梯度都是单独计算的。 因此完全可以把网络前向传播得到结果按之前的顺序拼接起来后再进行梯度的反向传播，这样就可以只进行一步训练，且不会出现训练先后的偏差。代码如下：

			f_out_y0, cf_out_y0 = net(x0)
            cf_out_y1, f_out_y1 = net(x1)

            #按照t=0和t=1的索引拼接向量
            y_pred = torch.zeros([len(x), 1])
            y_pred[index0] = f_out_y0
            y_pred[index1] = f_out_y1
            
			criterion = Loss()
            loss = criterion(f_y, y_pred, w) + 0.01 * (l2_regularization0 + l2_regularization1)
            #print(loss.item())
            viz.line([float(loss)], [epoch], win='train_loss', update='append')
            optimizer.zero_grad()
            loss.backward()

            #对网络的参数进行更新
            optimizer.step()

总结

two-head网络前向传播得到结果的时候是分开得到的，训练的时候通过拼接预测结果可以实现一次训练。