PyTorch有状态自定义损失函数中numpy操作的可行性及实现方案问询

我在PyTorch中实现了如下有状态自定义损失函数：

class PositionLossNormalized(torch.nn.Module):
def __init__(self, beta, epsilon=1e-8):
    super(PositionLossNormalized, self).__init__()
    self.beta = beta
    self.beta_t = beta
    self.running_avg = 0.0
    self.epsilon = epsilon

def forward(self, targets, predictions):
    sqr_diff = torch.square(targets - predictions)
    Pxj = torch.mean(sqr_diff, dim=-1)
    Pxj2 = torch.square(Pxj)
    
    avg_Pxj = torch.mean(Pxj)
    avg_Pxj2 = torch.mean(Pxj2)
    
    self.running_avg = self.beta * self.running_avg + (1.0 - self.beta) * avg_Pxj2.item()
    loss_val = avg_Pxj / (np.sqrt(self.running_avg / (1.0 - self.beta_t)) + self.epsilon)
    
    self.beta_t = self.beta * self.beta_t
    
    return loss_val

考虑到其中使用了numpy操作（即np.sqrt），请问该实现是否可行？我的思路是仅用Python标量对可自动求导的avg_Pxj进行缩放，不会破坏自动求导链，且运行代码时PyTorch未报错。我的理解是否正确？还是应该改用如下纯PyTorch函数实现？

class PositionLossNormalized(torch.nn.Module):
def __init__(self, beta, epsilon=1e-8):
    super(PositionLossNormalized, self).__init__()
    self.beta = torch.as_tensor(beta, dtype=torch.float32)
    self.beta_t = torch.as_tensor(beta, dtype=torch.float32)
    self.running_avg = torch.as_tensor(0.0, dtype=torch.float32)
    self.epsilon = torch.as_tensor(epsilon, dtype=torch.float32)

def forward(self, targets, predictions):
    sqr_diff = torch.square(targets - predictions)
    Pxj = torch.mean(sqr_diff, dim=-1)
    Pxj2 = torch.square(Pxj)
    
    avg_Pxj = torch.mean(Pxj)
    avg_Pxj2 = torch.mean(Pxj2)
    
    self.running_avg = self.beta * self.running_avg + (1.0 - self.beta) * avg_Pxj2.clone().detach()
    loss_val = avg_Pxj / (torch.sqrt(self.running_avg / (1.0 - self.beta_t)) + self.epsilon)
    
    self.beta_t = self.beta * self.beta_t
    
    return loss_val

def reset(self):
    self.beta_t = self.beta

解答 #

• 第一个实现确实可行，你的理解也正确：
  你用np.sqrt处理的是Python标量（running_avg、beta_t都是普通数值），对可导的avg_Pxj做的是标量除法操作——PyTorch会将这些标量视为常数，不会破坏avg_Pxj对应的计算图，求导时只会正常追踪avg_Pxj的梯度，因此运行时不会报错，逻辑上是通顺的。

• 更推荐使用第二个纯PyTorch实现，原因如下：

  1. 设备兼容性更好：如果模型部署在GPU上，第一个实现的numpy操作会强制将数据拉回CPU，虽然标量的性能影响可以忽略，但存在潜在的设备不匹配风险；纯PyTorch实现的张量会自动跟随模型所在设备，适配性更强。
  2. 风格统一易维护：全程使用PyTorch操作，避免混用numpy和PyTorch的数值逻辑，减少后续调试、扩展时的混淆。
  3. 精度与逻辑更一致：第一个实现中beta_t是Python浮点数，多次相乘会累积精度损失；PyTorch张量的数值计算逻辑和模型训练的精度控制更统一。
  4. 扩展性更强：如果后续需要适配多GPU并行、状态同步等场景，纯PyTorch的张量形式更容易修改调整。
     另外，第二个实现中用clone().detach()切断avg_Pxj2的梯度追踪，和第一个实现用.item()提取标量的效果完全一致，都是确保滑动平均状态running_avg不参与梯度计算，这部分处理是正确的。

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