PyTorch Dataloader批量维度异常导致LSTM输入报错问题排查

问题根源 #

你的自定义Dataset核心逻辑错误——__getitem__方法没有根据索引idx提取单个样本，而是直接返回整个数据集和标签数组。Dataloader会把64份完整数据集堆叠，导致最终inputs维度变成[64,10000,450,28]，而LSTM要求输入为3D张量（格式为[batch_size, seq_len, input_size]或[seq_len, batch_size, input_size]，由batch_first参数决定）。

修正后的Dataset代码 #

import torch
from torch.utils.data import Dataset

class CustomDataset(Dataset):
    def __init__(self, data, labels):
        # data维度：[样本数, 序列长度, 特征数] = [10000,450,28]
        # labels维度：[样本数, 序列长度] = [10000,450]
        self.data = torch.tensor(data, dtype=torch.float32)
        self.labels = torch.tensor(labels, dtype=torch.long)  # 分类任务标签用long类型

    def __len__(self):
        # 返回总样本数，即data的第一维长度
        return self.data.shape[0]

    def __getitem__(self, idx):
        # 根据索引取单个样本：特征维度[450,28]，标签维度[450]
        return self.data[idx], self.labels[idx]

Dataloader与LSTM适配示例 #

修正后构建Dataloader，遍历得到的inputs维度为[64,450,28]（batch_size=64，batch_first默认True），完全符合LSTM输入要求：

from torch.utils.data import DataLoader

# 假设你的numpy数组为np_data和np_labels
dataloader = DataLoader(CustomDataset(np_data, np_labels), batch_size=64, shuffle=True)

# 定义LSTM：输入特征数28，隐藏层维度64，返回所有时间步输出
lstm = torch.nn.LSTM(input_size=28, hidden_size=64, batch_first=True)

for inputs, labels in dataloader:
    # inputs维度：[64,450,28]
    # labels维度：[64,450]
    outputs, (h_n, c_n) = lstm(inputs)
    # outputs维度：[64,450,64]，可接全连接层实现逐帧分类

额外注意事项 #

• 必须将numpy数组转为PyTorch张量，并指定正确数据类型（特征用float32，分类标签用long）
• 逐帧分类时，LSTM输出的每个时间步结果对应一个标签，后续可通过Linear(64, num_classes)映射到目标分类类别

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