PyTorch DataLoader报TypeError：default_collate遇object类型数据求解

我在测试训练集批次处理，训练集是含3列的.tsv文件：Quality（1表示句子相似，0表示相反）、#1 String、#2 String。尝试把X、y转成列表等类型仍触发TypeError，后发现X的dtype为Object，合并X与y后整体dtype也为Object，求解决办法。

代码及错误信息 #

def get_dataloaders(ds, lengths=[0.6, 0.2, 0.2], batch_size=32, seed=42, num_workers=2):
    train_set, val_set, test_set = random_split(ds, lengths=lengths, generator=torch.Generator().manual_seed(seed))

    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=num_workers)
    val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=num_workers)
    test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=num_workers)
    return train_loader, val_loader, test_loader

data_dir = "/data.tsv"
data = pd.read_csv(data_dir, sep='\t')
y = data['Quality'].values                             #dtype: int64
X = data[['#1 String', '#2 String']].values            #dtype: O

data_input = np.column_stack((X, y))                   #dtype: O

train_loader, val_loader, test_loader = get_dataloaders(data_input)

for batch in train_loader:     #TypeError: default_collate: batch must contain tensors, numpy arrays, numbers, dicts or lists; found object

错误追踪信息：

TypeError: Caught TypeError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/_utils/worker.py", line 308, in _worker_loop
    data = fetcher.fetch(index)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/_utils/fetch.py", line 54, in fetch
    return self.collate_fn(data)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/_utils/collate.py", line 265, in default_collate
    return collate(batch, collate_fn_map=default_collate_fn_map)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/_utils/collate.py", line 119, in collate
    return collate_fn_map[elem_type](batch, collate_fn_map=collate_fn_map)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/_utils/collate.py", line 169, in collate_numpy_array_fn
    raise TypeError(default_collate_err_msg_format.format(elem.dtype))
TypeError: default_collate: batch must contain tensors, numpy arrays, numbers, dicts or lists; found object

解决方法 #

问题核心是直接把混合字符串和数值的numpy object数组传给DataLoader，PyTorch默认的collate函数无法处理这种混合类型。以下是几种可行方案：

方案1：自定义Dataset类（推荐） #

这是PyTorch处理文本数据的标准方式，拆分文本对和标签，避免混合类型数组问题：

import torch
from torch.utils.data import Dataset, DataLoader, random_split
import pandas as pd

class TextPairDataset(Dataset):
    def __init__(self, df):
        self.texts1 = df['#1 String'].tolist()
        self.texts2 = df['#2 String'].tolist()
        self.labels = df['Quality'].tolist()
    
    def __len__(self):
        return len(self.labels)
    
    def __getitem__(self, idx):
        # 返回文本对和标签，后续可在collate_fn或预处理步骤做tokenization
        return self.texts1[idx], self.texts2[idx], self.labels[idx]

# 加载数据并创建Dataset
data_dir = "/data.tsv"
data = pd.read_csv(data_dir, sep='\t')
dataset = TextPairDataset(data)

# 拆分数据集并生成DataLoader
train_set, val_set, test_set = random_split(dataset, lengths=[0.6, 0.2, 0.2], generator=torch.Generator().manual_seed(42))
train_loader = DataLoader(train_set, batch_size=32, shuffle=True, num_workers=2)
val_loader = DataLoader(val_set, batch_size=32, shuffle=False, num_workers=2)
test_loader = DataLoader(test_set, batch_size=32, shuffle=False, num_workers=2)

# 测试批次加载
for texts1, texts2, labels in train_loader:
    print(f"批次文本数量：{len(texts1)}，标签数量：{len(labels)}")
    break

方案2：自定义collate_fn #

如果坚持用numpy数组，给DataLoader传入自定义collate函数，手动处理每个样本：

import torch
import numpy as np
import pandas as pd
from torch.utils.data import DataLoader, random_split

def custom_collate(batch):
    # batch是列表，每个元素是numpy数组的一行（[str1, str2, label]）
    texts1 = [item[0] for item in batch]
    texts2 = [item[1] for item in batch]
    labels = torch.tensor([item[2] for item in batch], dtype=torch.long)
    return texts1, texts2, labels

def get_dataloaders(ds, lengths=[0.6, 0.2, 0.2], batch_size=32, seed=42, num_workers=2):
    train_set, val_set, test_set = random_split(ds, lengths=lengths, generator=torch.Generator().manual_seed(seed))

    # 添加自定义collate_fn参数
    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=num_workers, collate_fn=custom_collate)
    val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=num_workers, collate_fn=custom_collate)
    test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=num_workers, collate_fn=custom_collate)
    return train_loader, val_loader, test_loader

# 加载数据部分不变
data_dir = "/data.tsv"
data = pd.read_csv(data_dir, sep='\t')
y = data['Quality'].values
X = data[['#1 String', '#2 String']].values
data_input = np.column_stack((X, y))

train_loader, val_loader, test_loader = get_dataloaders(data_input)

# 测试批次
for texts1, texts2, labels in train_loader:
    print(f"批次文本数量：{len(texts1)}，标签形状：{labels.shape}")
    break

方案3：直接用DataFrame拆分 #

无需转成numpy数组，直接基于DataFrame创建Dataset：

class TextPairDataset(Dataset):
    def __init__(self, df):
        self.df = df
    
    def __len__(self):
        return len(self.df)
    
    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        return row['#1 String'], row['#2 String'], row['Quality']

data_dir = "/data.tsv"
data = pd.read_csv(data_dir, sep='\t')
dataset = TextPairDataset(data)

# 后续拆分和DataLoader创建同方案1

内容的提问来源于stack exchange，提问作者A Random Guy