使用BCEWithLogitsLoss训练后，如何用PyTorch Lightning实现预测

一、添加predict_step方法 #

在你的MUTAGClassifier类中加入predict_step方法，该方法会在调用trainer.predict()时自动处理每个batch，返回转换为0-1区间概率的预测值，以及对应的真实标签：

def predict_step(self, batch, batch_idx):
    # 前向传播获取模型输出的logits
    y_hat_logits = self.forward(batch.x, batch.edge_index, batch.batch)
    # 将logits转换为概率值
    y_hat_prob = torch.sigmoid(y_hat_logits).squeeze(1)
    # 获取真实标签并统一格式
    y_true = batch.y.float().squeeze(1)
    # 返回当前batch的预测与真实值
    return {"preds": y_hat_prob, "targets": y_true}

二、完整修改后的模型类代码 #

import torch
import torch.nn.functional as F
from sklearn.model_selection import train_test_split
import pytorch_lightning as ptlight
import torch_geometric as ptgeom
from torchmetrics.functional import accuracy

class MUTAGClassifier(ptlight.LightningModule):
    
    def __init__(self):
        # 模型结构：GCNConv -> GCNConv -> 图池化 -> Dropout -> Linear
        super().__init__()
        self.gc1 = ptgeom.nn.GCNConv(7, 256)
        self.gc2 = ptgeom.nn.GCNConv(256, 256)
        self.linear = torch.nn.Linear(256, 1)

    def forward(self, x, edge_index=None, batch=None, edge_weight=None):
        # 兼容单样本直接传入Data对象的场景
        if edge_index is None:
            x, edge_index, batch = x.x, x.edge_index, x.batch
        x = F.relu(self.gc1(x, edge_index, edge_weight))
        x = F.relu(self.gc2(x, edge_index, edge_weight))
        x = ptgeom.nn.global_mean_pool(x, batch)
        x = F.dropout(x, training=self.training)  # 仅训练阶段启用dropout
        x = self.linear(x)
        return x
 
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

    def training_step(self, batch, _):
        y_hat = self.forward(batch.x, batch.edge_index, batch.batch)
        loss = F.binary_cross_entropy_with_logits(y_hat, batch.y.unsqueeze(1).float())
        self.log("train_loss", loss)
        self.log("train_accuracy", accuracy(y_hat, batch.y.unsqueeze(1)), prog_bar=True, batch_size=32)
        return loss

    def validation_step(self, batch, _):
        x, edge_index, batch_idx = batch.x, batch.edge_index, batch.batch
        y_hat = self.forward(x, edge_index, batch_idx)
        self.log("val_accuracy", accuracy(y_hat, batch.y.unsqueeze(1)), prog_bar=True, batch_size=32)

    def predict_step(self, batch, batch_idx):
        y_hat_logits = self.forward(batch.x, batch.edge_index, batch.batch)
        y_hat_prob = torch.sigmoid(y_hat_logits).squeeze(1)
        y_true = batch.y.float().squeeze(1)
        return {"preds": y_hat_prob, "targets": y_true}

三、批量预测与结果收集 #

训练完成后，通过trainer.predict()获取所有测试样本的预测值和真实值，并整理为列表：

# 若从checkpoint加载模型，替换为：
# gnn = MUTAGClassifier.load_from_checkpoint("./checkpoints/gnn-200.ckpt")

# 执行预测
predictions = trainer.predict(gnn, dataloaders=test_loader)

# 合并所有batch的结果
all_preds = []
all_targets = []
for batch_pred in predictions:
    all_preds.extend(batch_pred["preds"].cpu().numpy().tolist())
    all_targets.extend(batch_pred["targets"].cpu().numpy().tolist())

# 查看前5个样本的结果
print("预测值列表（前5个）：", [round(p,4) for p in all_preds[:5]])
print("真实值列表（前5个）：", all_targets[:5])

四、单样本预测实现 #

直接从数据集中取单个样本，传入模型即可得到预测结果：

# 取测试集中的第一个样本
single_sample = mutag[test_idx[0]]

# 手动调用时切换到评估模式（trainer.predict会自动处理）
gnn.eval()
with torch.no_grad():
    # 直接传入样本对象，forward方法会自动解析x、edge_index
    logits = gnn(single_sample)
    pred_prob = torch.sigmoid(logits).item()
    true_label = single_sample.y.item()

print(f"单样本预测概率：{pred_prob:.4f}")
print(f"单样本真实标签：{true_label}")

关键说明 #

• predict_step中用sigmoid将模型输出的logits转换为概率值，符合二分类任务的预测习惯；
• forward方法兼容两种输入方式：既可以传入拆分的x、edge_index，也可以直接传入PyTorch Geometric的Data对象，方便单样本预测；
• 调用trainer.predict()时，模型会自动切换到评估模式，关闭dropout等训练专属层，保证预测结果稳定。

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