如何在PyTorch中将损失函数转化为得分函数实现得分最大化

问题背景 #

我有一个PyTorch模型，包含LSTM输入层、线性隐藏层，输出层为3个神经元并采用softmax激活。不想使用nn.MSELoss()这类依赖单样本目标标签的传统损失函数，而是要通过自定义得分函数score = get_score(predictions)实现最大化模型得分的训练目标。这个得分由模型整体行为计算而来，关联预测准确率和操作次数（比如预测0/1算操作，2算不操作，需平衡高准确率与合理操作率）。但训练时经常在loss.backward()阶段因梯度无法追踪报错，也不确定用单个数值得分评估全量预测是否合理。

简化示例代码 #

import torch
import torch.nn as nn
import torch.nn.functional as F
import random

class Model_123(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(10, 3)

    def forward(self, x):
        out = self.linear1(x)
        out = F.softmax(out, dim=1)
        return out 

dataset = torch.randn(10, 10)
model = Model_123()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
epochs = 100

custom_loss_score = CustomLoss_Score()  # 需要实现的类
for epoch in range(epochs):
    model.train()
    optimizer.zero_grad()
    predictions = model(dataset)
    
    # 示例：实际场景中score会基于predictions计算，这里用随机值占位
    # 真实逻辑：用predictions的argmax得到动作，统计动作0/1的成功率和操作次数，计算综合得分
    score = random.uniform(0, 100)

    loss_score = custom_loss_score(predictions, score)  # 目标是最大化score
    loss_score.backward()
    optimizer.step()

核心问题解决 #

1. 全量得分评估的合理性 #

全量评估是可行的，但得分函数必须规避“投机取巧”场景（比如只预测1个有效动作就拿100%准确率）。可以参考F1-score思路，设计同时兼顾成功率和操作覆盖率的得分公式，示例如下：

def calculate_score(predictions, true_actions):
    # true_actions为真实动作标签，无监督场景需替换为环境反馈或自洽性逻辑
    pred_actions = torch.argmax(predictions, dim=1)
    # 统计操作次数（动作0/1）
    action_count = torch.sum((pred_actions == 0) | (pred_actions == 1))
    operation_rate = action_count / len(pred_actions)
    # 统计成功次数（仅动作0/1预测正确算成功）
    success_count = torch.sum((pred_actions == true_actions) & ((pred_actions ==0)|(pred_actions==1)))
    success_rate = success_count / action_count if action_count >0 else 0.0
    # 综合得分：用F1-score平衡成功率与操作率
    score = 2 * (success_rate * operation_rate) / (success_rate + operation_rate + 1e-8)
    return score

2. 梯度可追踪的损失函数实现 #

PyTorch优化器默认最小化损失，因此要最大化得分，只需将损失定义为-score。关键是保证得分计算全程基于可微分张量操作，避免argmax这类离散操作（会切断梯度流）。

解决方案：用软标签替代硬分类

直接用softmax输出的概率计算得分，避免离散化操作：

class CustomLoss_Score(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, predictions, true_actions=None):
        # 计算操作概率：动作0和1的概率之和
        operation_probs = predictions[:,0] + predictions[:,1]
        operation_rate = torch.mean(operation_probs)
        
        # 监督场景：计算成功概率
        if true_actions is not None:
            # 取每个样本对应真实动作的概率
            correct_probs = torch.gather(predictions, 1, true_actions.unsqueeze(1)).squeeze()
            # 仅统计动作0/1的样本
            mask = (true_actions ==0) | (true_actions ==1)
            success_probs = correct_probs[mask]
            success_rate = torch.mean(success_probs) if len(success_probs)>0 else 0.0
        
            # 综合得分（F1变体）
            score = 2 * (success_rate * operation_rate) / (success_rate + operation_rate + 1e-8)
        else:
            # 无监督场景示例：鼓励模型输出多样化有效动作
            score = torch.mean(operation_probs) - torch.var(predictions[:,0:1])
        
        # 最大化score等价于最小化-score
        return -score

修改后的训练循环

# 监督场景示例：假设有真实动作标签，无监督场景去掉true_actions即可
true_actions = torch.randint(0,3,(10,))
custom_loss_score = CustomLoss_Score()

for epoch in range(epochs):
    model.train()
    optimizer.zero_grad()
    predictions = model(dataset)
    
    # 传入predictions和true_actions（无监督场景仅传predictions）
    loss_score = custom_loss_score(predictions, true_actions)
    loss_score.backward()
    optimizer.step()
    
    # 打印训练日志
    if epoch %10 ==0:
        current_score = -loss_score.item()
        print(f"Epoch {epoch}: Score = {current_score:.4f}")

3. 梯度中断问题的关键修复 #

• 避免所有离散操作（如argmax、张量转整数、numpy转换），所有计算基于PyTorch可微分张量操作。
• 若必须使用离散动作（如与环境交互），可采用Gumbel-Softmax技巧：训练时用近似可微分的软采样，测试时用硬分类：

def gumbel_softmax(logits, tau=1.0, hard=False):
    return F.gumbel_softmax(logits, tau=tau, hard=hard, dim=1)

# 修改模型forward方法
def forward(self, x):
    logits = self.linear1(x)
    # 训练用软采样，测试用硬采样
    if self.training:
        out = gumbel_softmax(logits, tau=0.5, hard=False)
    else:
        out = F.softmax(logits, dim=1)
    return out

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