训练Vision Transformer时np.split等分失败的ValueError问题

报错代码行：

x = np.split(np.squeeze(np.array(x)), BATCH_SIZE)

已尝试遍历候选批量大小，选取能整除训练样本数的BATCH_SIZE，但训练启动后仍触发错误，相关代码如下：

训练核心代码 #

import torch.utils.data as data
from torch.autograd import Variable
import numpy as np

print("Number of train samples: ", len(train_ds))
print("Number of test samples: ", len(test_ds))
print("Detected Classes are: ", train_ds.class_to_idx) 

train_loader = data.DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True,  num_workers=4)
test_loader  = data.DataLoader(test_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=4) 

# Train the model
for epoch in range(EPOCHS):        
  for step, (x, y) in enumerate(train_loader):
    # Change input array into list with each batch being one element
    x = np.split(np.squeeze(np.array(x)), BATCH_SIZE)
    # Remove unecessary dimension
    for index, array in enumerate(x):
      x[index] = np.squeeze(array)
    # Apply feature extractor, stack back into 1 tensor and then convert to tensor
    x = torch.tensor(np.stack(feature_extractor(x)['pixel_values'], axis=0))
    # Send to GPU if available
    x, y  = x.to(device), y.to(device)
    b_x = Variable(x)   # batch x (image)
    b_y = Variable(y)   # batch y (target)
    # Feed through model
    output, loss = model(b_x, None)
    # Calculate loss
    if loss is None: 
      loss = loss_func(output, b_y)   
      optimizer.zero_grad()           
      loss.backward()                 
      optimizer.step()

    if step % 50 == 0:
      # Get the next batch for testing purposes
      test = next(iter(test_loader))
      test_x = test[0]
      # Reshape and get feature matrices as needed
      test_x = np.split(np.squeeze(np.array(test_x)), BATCH_SIZE)
      for index, array in enumerate(test_x):
        test_x[index] = np.squeeze(array)
      test_x = torch.tensor(np.stack(feature_extractor(test_x)['pixel_values'], axis=0))
      # Send to appropirate computing device
      test_x = test_x.to(device)
      test_y = test[1].to(device)
      # Get output (+ respective class) and compare to target
      test_output, loss = model(test_x, test_y)
      test_output = test_output.argmax(1)
      # Calculate Accuracy
      accuracy = (test_output == test_y).sum().item() / BATCH_SIZE
      print('Epoch: ', epoch, '| train loss: %.4f' % loss, '| test accuracy: %.2f' % accuracy)

批量大小选择代码 #

import math

# Determine batch size
batch_size_candidates = [32, 64, 128, 256]
for candidate in batch_size_candidates:
    if len(train_ds) % candidate == 0:
        BATCH_SIZE = candidate
        break
if BATCH_SIZE is None:
    raise ValueError("Cannot find a batch size that evenly divides the number of training samples.")
print("Batch size: ", BATCH_SIZE)

问题原因分析 #

1. DataLoader默认保留不完整批次：即使训练集样本数能被BATCH_SIZE整除，测试集样本数可能不满足；且DataLoader默认不会丢弃最后一个样本数不足的批次，当处理这类批次时，用固定的BATCH_SIZE做split会触发分割不均的错误。
2. np.squeeze破坏维度结构：如果输入张量存在维度为1的轴（比如单通道图像的通道维度），np.squeeze会盲目删除该维度，导致数组形状不符合预期，无法按BATCH_SIZE均分。
3. 测试逻辑存在漏洞：用next(iter(test_loader))获取测试批次时，同样可能拿到不完整批次；且计算准确率时固定用BATCH_SIZE做分母，会导致准确率计算失真。

解决方案 #

1. 配置DataLoader丢弃不完整批次 #

创建DataLoader时添加drop_last=True参数，确保每个批次的样本数严格等于BATCH_SIZE：

train_loader = data.DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=4, drop_last=True)
test_loader  = data.DataLoader(test_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=4, drop_last=True)

2. 简化输入处理，避免手动split #

DataLoader已返回按批次组织的张量，无需手动split，直接遍历样本处理即可，同时避免盲目挤压维度：

# 训练循环内的输入处理替换为：
x_np = np.array(x)
pixel_values_list = []
for idx in range(x_np.shape[0]):
    sample = x_np[idx]
    # 仅删除指定的冗余维度（比如单通道图像的通道轴）
    if sample.ndim == 3 and sample.shape[0] == 1:
        sample = sample.squeeze(0)
    feat = feature_extractor(sample)['pixel_values']
    pixel_values_list.append(feat)
x = torch.tensor(np.stack(pixel_values_list, axis=0))

3. 修正测试部分逻辑 #

避免重复创建测试迭代器，同时根据实际批次大小计算准确率：

if step % 50 == 0:
    test_x, test_y = next(iter(test_loader))
    # 处理测试输入
    test_np = np.array(test_x)
    test_pixel_values = []
    for idx in range(test_np.shape[0]):
        sample = test_np[idx]
        if sample.ndim == 3 and sample.shape[0] == 1:
            sample = sample.squeeze(0)
        feat = feature_extractor(sample)['pixel_values']
        test_pixel_values.append(feat)
    test_x = torch.tensor(np.stack(test_pixel_values, axis=0)).to(device)
    test_y = test_y.to(device)
    
    test_output, loss = model(test_x, test_y)
    test_output = test_output.argmax(1)
    # 用实际批次大小计算准确率
    batch_acc = (test_output == test_y).sum().item() / test_y.size(0)
    print('Epoch: ', epoch, '| train loss: %.4f' % loss, '| test accuracy: %.2f' % batch_acc)

4. 兜底方案：动态获取当前批次大小 #

如果不想丢弃最后一个批次，可在split时用当前批次的实际样本数替代BATCH_SIZE：

# 训练循环内：
current_batch_size = x.size(0)
x = np.split(np.squeeze(np.array(x)), current_batch_size)

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