如何正确切片PyTorch张量以获取[2,64]规格的批次数据？

问题背景 #

我有一个形状为torch.Size([2, 11938])的PyTorch张量，想要切分为[2,64]规格的批次用于训练，现有代码如下：

batch = 0
BATCH_SIZE = 64
X_train = x_scaled.to(device)
y_train = y_scaled.to(device)

for batch in range(0,len(X_train[0]),BATCH_SIZE):
   ### Training
   model.train() # train mode is on by default after construction
   # 1. Forward pass
   y_pred = model(X_train[0:2][batch:batch+BATCH_SIZE])

运行后出现错误：mat1 and mat2 shapes cannot be multiplied (2x11938 and 2x64)，且切片结果不符合预期：

• 期望得到形状为[2,64]的张量，示例如下：

tensor([[0.0000, 0.0002, 0.0004, 0.0005, 0.0007, 0.0009, 0.0011, 0.0013, 0.0014,
    0.0016, 0.0018, 0.0018, 0.0020, 0.0022, 0.0023, 0.0025, 0.0027, 0.0029,
    0.0029, 0.0031, 0.0032, 0.0034, 0.0036, 0.0038, 0.0040, 0.0041, 0.0043,
    0.0045, 0.0047, 0.0049, 0.0051, 0.0052, 0.0054, 0.0056, 0.0058, 0.0060,
    0.0061, 0.0061, 0.0063, 0.0065, 0.0067, 0.0069, 0.0070, 0.0072, 0.0074,
    0.0076, 0.0078, 0.0079, 0.0081, 0.0083, 0.0083, 0.0085, 0.0087, 0.0088,
    0.0090, 0.0092, 0.0094, 0.0094, 0.0096, 0.0097, 0.0099, 0.0101, 0.0103,
    0.0105],[0.0684, 0.0684, 0.0684, 0.0684, 0.0684, 0.0684, 0.0684, 0.0684, 0.0684,
    0.0703, 0.0703, 0.0703, 0.0684, 0.0684, 0.0703, 0.0703, 0.0703, 0.0703,
    0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703,
    0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703,
    0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703,
    0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703,
    0.0703, 0.0703, 0.0712, 0.0712, 0.0712, 0.0712, 0.0712, 0.0712, 0.0712,
    0.0712]], device='cuda:0', dtype=torch.float64)

• 实际得到的张量仍保持[2,11938]形状，示例如下：

tensor([[0.0000e+00, 1.8038e-04, 3.6076e-04,  ..., 9.9964e-01, 9.9982e-01,
         1.0000e+00],
        [6.8395e-02, 6.8395e-02, 6.8395e-02,  ..., 5.7695e-01, 5.7695e-01,
         5.7695e-01]], device='cuda:0', dtype=torch.float64)

错误原因 #

你的切片逻辑有误：X_train[0:2][batch:batch+BATCH_SIZE]是先取前2行（张量本身只有2行，等于没切片），再对行维度进行切片，但你需要的是对列维度进行切片，同时保留所有行。

解决方案 #

方法1：修正切片语法 #

直接在张量的第二个维度（列维度）上进行切片，语法为X_train[:, batch:batch+BATCH_SIZE]，完整代码如下：

BATCH_SIZE = 64
X_train = x_scaled.to(device)
y_train = y_scaled.to(device)

for batch_idx in range(0, X_train.shape[1], BATCH_SIZE):
    model.train()
    # 对列维度切片，保留所有行，得到形状为[2,64]的批次
    X_batch = X_train[:, batch_idx:batch_idx+BATCH_SIZE]
    y_pred = model(X_batch)
    # 后续训练步骤（损失计算、反向传播等）

• X_train.shape[1]直接获取列维度长度，比len(X_train[0])更直观
• 循环变量改为batch_idx避免变量名混淆

方法2：使用torch.split直接分割 #

如果不想手动写循环切片，可以用torch.split一次性按指定大小分割张量：

BATCH_SIZE = 64
X_train = x_scaled.to(device)
y_train = y_scaled.to(device)

# 按列维度分割为多个[2,64]或[2,剩余长度]的张量
batches = torch.split(X_train, BATCH_SIZE, dim=1)
for X_batch in batches:
    model.train()
    y_pred = model(X_batch)
    # 后续训练步骤

注意：最后一个批次的长度可能小于64（11938 % 64 = 34），如果模型要求固定输入长度，可以根据需求补零或丢弃该批次。

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