PyTorch模型输出到预测标签转换函数Logits2PredLabel

函数Logits2PredLabel

def Logits2PredLabel(logits,model_out_len):
    'pred -> hard label'
    with torch.no_grad():
        if model_out_len == 1:
            prediction = torch.round(torch.sigmoid(logits))
        else:
            _,prediction = torch.max(logits,1)
    return prediction

功能说明

这个函数的作用是将模型的输出转换为预测的标签。对于只有一个输出的模型，使用sigmoid函数将logits转换为0或1的预测值。对于有多个输出的模型，选择具有最大值的索引作为预测标签。函数返回预测的标签。