双神经网络交叉正则化失效问题排查及代码求助

我尝试为两个神经网络添加损失项实现交叉正则化，让它们在完成不同任务的同时尽可能相似。参考相关方案后尝试实现，但模型训练后两个网络精度不错，却完全无视正则化约束——哪怕把alpha设得极高，两网络的差异反而持续增大。请问代码哪里出错了？需要对正则化损失项做哪些处理才能避免被忽略？

用户原尝试代码：

def train_combined(nets, dataset_train, dataset_test, num_epochs, alpha=0):
  criterion = nn.L1Loss()
  optimizers = [optim.SGD(net.parameters(), lr=0.01, momentum=0.9 ) for net in nets]
  trainloader = DataLoader(dataset_train, batch_size=32, shuffle=True )

  train_losses = []
  test_losses  = []

  for epoch in range(num_epochs):  # loop over the dataset multiple times

    for i, data in enumerate(trainloader, 0):
      # get the inputs; data is a list of [inputs, labels]
      inputs, *labels = data
      inputs = inputs

      # get the average of the paramaters between the two networks
      with t.no_grad():
        params = t.stack([t.cat(tuple(t.flatten(p.data) for p in net.parameters())) for net in nets])
        avg = t.sum(params, dim=0)*0.5

      # keep track of loss for both models
      all_losses = np.zeros( 2 )
      all_reg_losses = np.zeros( 2 )
      all_final_losses = np.zeros( 2 )

      # forward + backward + optimize
      for i, (net, optimizer, label) in enumerate(zip(nets, optimizers, labels)):
        optimizer.zero_grad()

        # calculate normal loss
        outputs = net(inputs)
        loss = criterion(outputs, label)        
        
        # calculate regularization loss loss
        params = t.cat(tuple(t.flatten(p.data) for p in net.parameters()))
        regularization_loss = t.sum(t.abs( params - avg ))
        regularization =  regularization_loss * alpha
        
        # calculate total loss
        final_loss = loss + regularization
        final_loss.backward()
        optimizer.step()

        # keep track of losses
        all_losses[i] = float( loss.item() )
        all_reg_losses[i] = float( 0 if (regularization == 0) else regularization.item() )
        all_final_losses[i] = float( final_loss.item() )
      
      # keep track of performance
      train_losses.append( loss )
      with t.no_grad():
        for i in range(2):
          test_losses.append( light_eval( nets[i], data_test, index=i ) )

    # log performance each epoch
    for i in range(2):
      print("%3d" % (epoch+1),  i, ':',
            f'  train loss  = { ("%.4f "*3) % (all_losses[i], all_reg_losses[i], all_final_losses[i]) }',
            f', test_losses = { "%.4f" % test_losses[-(2-i)] }')
    

  print('Finished Training')

models = [ Net().to(device) for i in range(2) ]
train_combined( models, dataset_train, dataset_test, 50, alpha=1e-2 )

• 正则化损失丢失梯度追踪：计算正则化损失时用了p.data，这会剥离张量的梯度信息，导致正则化项的损失无法反向传播，优化器完全不会对正则化约束做出响应。
• 参数平均时机错误：先计算全局参数平均值，再依次更新两个网络。更新第一个网络后，第二个网络的正则化约束仍基于初始的平均值，导致约束逻辑不一致，无法有效让参数靠拢。
• 正则化损失尺度未归一化：直接对所有参数的绝对值差求和，参数数量多的话损失值会远大于任务损失，alpha调节难度极大——要么正则化太弱，要么直接覆盖任务损失但优化逻辑混乱。
• 损失记录与测试逻辑混乱：train_losses只记录最后一个网络的任务损失；data_test未定义，test_losses的追加逻辑会导致日志索引错误。

import torch as t
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import numpy as np

def light_eval(net, dataset_test, index):
    testloader = DataLoader(dataset_test, batch_size=32, shuffle=False)
    criterion = nn.L1Loss()
    total_loss = 0.0
    with t.no_grad():
        for data in testloader:
            inputs, *labels = data
            inputs = inputs.to(next(net.parameters()).device)
            label = labels[index].to(next(net.parameters()).device)
            outputs = net(inputs)
            loss = criterion(outputs, label)
            total_loss += loss.item()
    return total_loss / len(testloader)

def train_combined(nets, dataset_train, dataset_test, num_epochs, alpha=0):
    criterion = nn.L1Loss()
    optimizers = [optim.SGD(net.parameters(), lr=0.01, momentum=0.9) for net in nets]
    trainloader = DataLoader(dataset_train, batch_size=32, shuffle=True)
    device = next(nets[0].parameters()).device

    # 分开记录两个网络的损失
    train_total_losses = [[], []]
    test_losses = [[], []]

    for epoch in range(num_epochs):
        running_task_loss = [0.0, 0.0]
        running_reg_loss = [0.0, 0.0]
        running_total_loss = [0.0, 0.0]

        for data in trainloader:
            inputs, *labels = data
            inputs = inputs.to(device)
            labels = [lbl.to(device) for lbl in labels]

            # 1. 同步获取两个网络的参数，保留梯度追踪
            params_list = []
            for net in nets:
                params = t.cat([t.flatten(p) for p in net.parameters()])
                params_list.append(params)
            params_stack = t.stack(params_list)
            avg_params = params_stack.mean(dim=0)

            # 2. 先计算两个网络的总损失，暂不更新参数
            all_final_losses = []
            for idx, (net, optimizer, label) in enumerate(zip(nets, optimizers, labels)):
                optimizer.zero_grad()
                outputs = net(inputs)
                task_loss = criterion(outputs, label)
                
                # 使用参数本身，保留梯度
                current_params = t.cat([t.flatten(p) for p in net.parameters()])
                # 用均值代替求和，归一化损失尺度
                reg_loss = t.mean(t.abs(current_params - avg_params))
                total_loss = task_loss + alpha * reg_loss
                
                all_final_losses.append(total_loss)
                running_task_loss[idx] += task_loss.item()
                running_reg_loss[idx] += reg_loss.item()
                running_total_loss[idx] += total_loss.item()

            # 3. 统一反向传播并更新参数，避免参数更新影响正则化约束
            for total_loss, optimizer in zip(all_final_losses, optimizers):
                total_loss.backward(retain_graph=True)
                optimizer.step()

        # 计算epoch平均损失
        epoch_task_loss = [loss / len(trainloader) for loss in running_task_loss]
        epoch_reg_loss = [loss / len(trainloader) for loss in running_reg_loss]
        epoch_total_loss = [loss / len(trainloader) for loss in running_total_loss]

        # 记录并打印日志
        for idx in range(2):
            train_total_losses[idx].append(epoch_total_loss[idx])
            test_loss = light_eval(nets[idx], dataset_test, idx)
            test_losses[idx].append(test_loss)
            
            print(f"{epoch+1:3d} 网络{idx} : "
                  f"训练损失 = 任务{epoch_task_loss[idx]:.4f} + 正则{epoch_reg_loss[idx]:.4f} = {epoch_total_loss[idx]:.4f}, "
                  f"测试损失 = {test_losses[idx][-1]:.4f}")

    print('训练完成')

# 示例网络，替换为你的实际网络结构
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 10)
    
    def forward(self, x):
        x = x.flatten(1)
        x = t.relu(self.fc1(x))
        return self.fc2(x)

device = t.device("cuda" if t.cuda.is_available() else "cpu")
models = [Net().to(device) for _ in range(2)]
# train_combined(models, dataset_train, dataset_test, 50, alpha=1.0)

• 保留梯度追踪：计算正则化损失时直接使用参数张量p而非p.data，确保正则化项的梯度能反向传播到参数。
• 同步计算损失再更新：先计算两个网络的总损失，再统一反向传播更新参数，避免更新第一个网络后破坏参数平均值的一致性。
• 正则化损失归一化：用t.mean()代替t.sum()计算绝对值差的均值，让正则化损失尺度与任务损失匹配，alpha参数调节更直观。
• 修正损失记录逻辑：分开记录两个网络的训练/测试损失，按epoch计算平均损失，日志输出更清晰准确。

内容的提问来源于stack exchange，提问作者Nicky Pochinkov