PyTorch从DataSet取单样本预测报错，用DataLoader正常如何解决？

我想从测试DataSet对象中随机抽取一个样本，使用训练好的模型进行预测，编写了如下代码但运行报错：

rng = np.random.default_rng()
ind = rng.integers(0,len(test_ds),(1,))[-1]


I = test_ds[ind] # Note I is a list of tensors of equal size
I = [Ik.to(device) for Ik in I]

with torch.no_grad():
      _, y_f_hat, _, y_f = model.forward_F(I)
      y_f_hat = y_f_hat.cpu().numpy().flatten()
      y_f = y_f.cpu().numpy().flatten()

报错信息：

ERROR: /usr/local/lib/python3.8/dist-packages/torch/nn/modules/flatten.py in forward(self, input)
     44 
     45     def forward(self, input: Tensor) -> Tensor:
---> 46         return input.flatten(self.start_dim, self.end_dim)
     47 
     48     def extra_repr(self) -> str:

IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)

但使用DataLoader遍历测试集时则完全正常，代码如下：

for I in test_dataloader:
  with torch.no_grad():
      _, y_f_hat, _, y_f = model.forward_F(I)
      y_f_hat = y_f_hat.cpu().numpy().flatten()
      y_f = y_f.cpu().numpy().flatten()
      break

注：test_ds是test_dataloader所使用的数据集，环境为Google Colab GPU、Python 3.9

问题根源是输入维度不匹配：

• DataLoader在迭代时会自动为每个样本添加batch维度（例如原本形状为[C, H, W]的张量，经过DataLoader后会变成[1, C, H, W]），符合模型forward_F的批量输入设计。
• 直接从DataSet取样本时，得到的是无batch维度的单样本张量，模型中的Flatten层（或其他依赖batch维度的层）无法处理，因此抛出维度越界错误。

修复代码只需给每个张量手动添加batch维度，修改后的代码如下：

rng = np.random.default_rng()
ind = rng.integers(0, len(test_ds), (1,))[-1]

I = test_ds[ind]
# 为每个张量添加batch维度（在第0位插入维度）
I = [Ik.unsqueeze(0).to(device) for Ik in I]

with torch.no_grad():
    _, y_f_hat, _, y_f = model.forward_F(I)
    y_f_hat = y_f_hat.cpu().numpy().flatten()
    y_f = y_f.cpu().numpy().flatten()

也可以用Ik[None, ...]替代Ik.unsqueeze(0)，两者效果一致，都是在张量最前面添加一个大小为1的batch维度。

这样处理后，输入张量的维度就和DataLoader输出的格式完全一致，模型即可正常处理单样本预测。

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