PyTorch模型推理时，如何输入无batch维度的单样本？

你的问题根源在于模型训练时接收的输入包含batch维度（第一个维度250是batch size），而推理时的单样本输入缺少这个维度，导致模型的层无法匹配输入形状。解决方法很简单，给单样本张量添加batch维度即可，具体有两种常用方式：

方法1：使用unsqueeze(0)动态添加batch维度 #

直接对单样本张量调用unsqueeze(0)，将形状从(120,8)转换为(1,120,8)（batch size=1），匹配训练时的输入维度结构：

# 正确的单样本推理代码
single_sample = torch.rand(120,8)
output = model(single_sample.unsqueeze(0), 'cpu')
# 如果需要去掉输出的batch维度（输出形状从(1, desired_output_size)变为(desired_output_size)）
output = output.squeeze(0)

方法2：直接创建带batch维度的单样本张量 #

在生成输入张量时就指定batch维度为1，省去后续维度转换步骤：

# 直接生成形状为(1,120,8)的单样本
output = model(torch.rand(1,120,8), 'cpu')

可选：让模型自动兼容两种输入形状 #

如果希望模型既能处理带batch维度的批量输入，又能直接接收单样本输入，可以修改模型的forward方法，添加维度检查和自动转换逻辑：

# 假设你的transformer.linear模型的forward方法可以修改
class LinearModel(nn.Module):
    def __init__(self, hparams):
        super().__init__()
        # 原有初始化逻辑
        self.num_features = hparams['num_features']
        self.seq_len = hparams['seq_len']
        self.linear = nn.Linear(self.seq_len * self.num_features, desired_output_size) # 示例层

    def forward(self, x, device):
        # 检查输入维度，如果是2维（无batch）则添加batch维度
        if x.dim() == 2:
            x = x.unsqueeze(0)
        # 原有forward逻辑，比如展平特征等
        x = x.flatten(1)
        x = self.linear(x)
        return x

修改后，无论是输入(250,120,8)还是(120,8)，模型都能正常处理。

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