PyTorch模型单样本与批量推理预测结果不一致问题咨询

您好，我发现用PyTorch训练的Transformer模型，在单样本推理和批量推理时，同一个样本的预测结果差异非常明显：单样本推理输出是[0.37732467 0.2642143 0.35846105]，但把这个样本放到批量里推理时，输出变成了[0.3185594 0.40971586 0.2717247 ]。想请教如何保证两种场景下的预测结果一致？

以下是我的完整代码：

from transformers import Trainer, TrainingArguments, PreTrainedModel, PretrainedConfig
from torch.utils.data import Dataset
import torch
import torch.nn.functional as F
import numpy as np

# Number of Features
num_of_features = 128

# Dataset Class
class SequenceDataset(Dataset):
    def __init__(self, X, y):
        self.X = torch.tensor(X, dtype=torch.float32)
        self.y = torch.tensor(y, dtype=torch.long)

    def __len__(self):
        return len(self.y)

    def __getitem__(self, idx):
        return {"input_ids": self.X[idx], "labels": self.y[idx]}


# Configuration Class
class SequenceConfig(PretrainedConfig):
    model_type = "sequence_transformer"

    def __init__(self, num_features=num_of_features, num_classes=3, d_model=1024, nhead=4, num_layers=4, dim_feedforward=512, **kwargs):
        self.num_features = num_features
        self.num_classes = num_classes
        self.d_model = d_model
        self.nhead = nhead
        self.num_layers = num_layers
        self.dim_feedforward = dim_feedforward
        super().__init__(**kwargs)


# Transformer Model
class SequenceTransformer(PreTrainedModel):
    config_class = SequenceConfig

    def __init__(self, config):
        super().__init__(config)
        self.embedding = torch.nn.Linear(config.num_features, config.d_model)
        self.positional_encoding = torch.nn.Parameter(torch.zeros(1, config.d_model))
        encoder_layer = torch.nn.TransformerEncoderLayer(
            d_model=config.d_model, 
            nhead=config.nhead, 
            dim_feedforward=config.dim_feedforward, 
            batch_first=True
        )
        self.transformer_encoder = torch.nn.TransformerEncoder(encoder_layer, num_layers=config.num_layers)
        self.fc = torch.nn.Linear(config.d_model, config.num_classes)

    def forward(self, input_ids, labels=None):
        src = self.embedding(input_ids) + self.positional_encoding
        output = self.transformer_encoder(src)
        logits = self.fc(output)
        probs = F.softmax(logits, dim=-1)

        loss = None
        if labels is not None:
            loss_fct = torch.nn.CrossEntropyLoss()
            loss = loss_fct(logits, labels)
            
        return {"loss": loss, "logits": logits, "probs": probs} if labels is not None else logits


# Training Code
config = SequenceConfig()
model = SequenceTransformer(config)

# Training Arguments
batchSize=32
numWarmUpSteps=int(np.shape(train_image)[0]/batchSize/numOfBreakpointsPerEpoch/10)
training_args = TrainingArguments(
    output_dir=path,
    num_train_epochs=1, 
    per_device_train_batch_size=batchSize,
    per_device_eval_batch_size=320,
    warmup_steps=numWarmUpSteps,
    weight_decay=0.1,
    logging_strategy='no',
    eval_strategy="epoch",
    save_strategy="epoch",
    metric_for_best_model="accuracy",
    save_only_model=True,
)

# Trainer Initialization
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    compute_metrics=compute_metrics
)

# Train the Model
train_output = trainer.train()

# Save Model and Training Arguments
trainer.save_model("./SavedModels")
torch.save(training_args, "./SavedModels/training_args.bin")

# Prediction Code
training_args_loaded = torch.load("./SavedModels/training_args.bin")
model_save_path = "./SavedModels/"
model = SequenceTransformer(config).from_pretrained(model_save_path)

trainer = Trainer(model=model, compute_metrics=compute_metrics, args=training_args_loaded)
test_data = np.random.rand(10, num_of_features)  # Example test data
test_predictions = trainer.predict(torch.tensor(test_data, dtype=torch.float32))

# Output Test Predictions
print(test_predictions)

问题分析与解决方案 #

这种单样本和批量结果不一致的情况，大概率是模型模式或数据处理的问题，下面是具体的排查和修复步骤：

• 强制切换模型到评估模式
  模型在训练时会启用dropout、动态层归一化等训练专属行为，推理时必须关闭这些才能保证结果稳定。你加载模型后没有显式调用model.eval()，虽然Trainer.predict()会自动切换模式，但如果自己手动做单样本推理时忘了切换，就会出现差异。
  建议在所有推理代码前加上：

  model.eval()
  

  同时搭配torch.no_grad()禁用梯度计算，既节省内存也避免干扰：

  with torch.no_grad():
      # 单样本推理示例
      single_sample = torch.tensor(test_data[0:1], dtype=torch.float32)
      single_logits = model(single_sample)
      single_probs = F.softmax(single_logits, dim=-1)
      print(single_probs.numpy())
  

• 确保数据预处理完全一致
  检查单样本和批量样本的预处理流程是否完全同步：比如数据类型是否都是float32、有没有遗漏的归一化步骤、输入形状是否匹配（单样本要保持(1, num_features)的批量维度，不能直接传(num_features)的张量）。你当前的测试数据是随机生成的，看起来没问题，但实际场景中要严格对齐。

• 验证Trainer的单样本预测行为
  可以直接用Trainer来做单样本预测，对比结果是否和批量一致：

  single_test_dataset = SequenceDataset(test_data[0:1], np.array([0]))
  single_pred_result = trainer.predict(single_test_dataset)
  print(single_pred_result.predictions)
  

  如果结果和批量一致，说明问题出在你自己手动单样本推理时的模式切换或数据处理上。

• 排查设备与精度问题
  偶尔GPU和CPU的浮点数计算精度差异，或者混合精度训练的残留影响会导致微小差异，但你这里的差异很大，所以这个可能性较低。可以尝试把模型和数据都放到同一个设备（比如统一用CPU或GPU）再测试。

备注：内容来源于stack exchange，提问作者Dishant Dua