PyTorch DataLoader训练遇UserWarning：列表转张量过慢问题求助

警告信息 #

/usr/local/lib/python3.10/dist-packages/sentence_transformers/SentenceTransformer.py:547:
UserWarning: Creating a tensor from a list of numpy.ndarrays is
extremely slow. Please consider converting the list to a single
numpy.ndarray with numpy.array() before converting to a tensor.
(Triggered internally at ../torch/csrc/utils/tensor_new.cpp:245.)
labels = torch.tensor(labels)

用户代码 #

from torch.utils.data import DataLoader
from sentence_transformers import losses
from sentence_transformers import ParallelSentencesDataset
from sentence_transformers import models
from sentence_transformers import SentenceTransformer

xlmr = models.Transformer('xlm-roberta-base')
pooler = models.Pooling(
    xlmr.get_word_embedding_dimension(),
    pooling_mode_mean_tokens=True
)

student = SentenceTransformer(modules=[xlmr, pooler])
teacher = SentenceTransformer('paraphrase-distilroberta-base-v2')

data = ParallelSentencesDataset(
    student_model=student,
    teacher_model=teacher,
    batch_size=32,
    use_embedding_cache=True
)
data.load_data('/path/to/somefile', max_sentence_length=512)

loader = DataLoader(data, shuffle=True, batch_size=32)
loss = losses.MSELoss(model=student)
  
epochs=1
student.fit(
    train_objectives=[(loader, loss)],
    epochs=epochs,
    warmup_steps=int(len(loader) * epochs * 0.1), # 10% of data
    output_path='./xlmr-ted',
    optimizer_params={'lr': 2e-5, 'eps': 1e-6},
    save_best_model=True,
    show_progress_bar=True
)

解决方案 #

警告根源是ParallelSentencesDataset生成的批量标签是numpy数组的列表，而非单个numpy数组，导致torch.tensor()转换效率低下。以下是几种可行方案：

方案1：自定义数据集的collate函数 #

通过包装原数据集，重写批量处理逻辑，将标签列表转换为单个numpy数组后再转张量：

import numpy as np
from torch.utils.data import Dataset

class NumpyArrayDataset(Dataset):
    def __init__(self, base_dataset):
        self.base_dataset = base_dataset
    
    def __len__(self):
        return len(self.base_dataset)
    
    def __getitem__(self, idx):
        return self.base_dataset[idx]
    
    def collate_fn(self, batch):
        # 拆分输入与标签
        inputs = [item[0] for item in batch]
        labels = [item[1] for item in batch]
        
        # 将标签列表转为单个numpy数组
        labels = np.array(labels)
        return inputs, labels

修改DataLoader的创建代码：

wrapped_data = NumpyArrayDataset(data)
loader = DataLoader(wrapped_data, shuffle=True, batch_size=32, collate_fn=wrapped_data.collate_fn)

方案2：预计算教师模型嵌入（更高效） #

利用use_embedding_cache=True的特性，提前计算所有教师模型的嵌入并保存为单个numpy数组，避免动态生成数组列表：

# 预计算所有句子的教师嵌入，直接得到numpy数组
teacher_embeddings = teacher.encode(data.sentences, batch_size=32, convert_to_numpy=True)
# 替换数据集的缓存嵌入
data.teacher_embeddings = teacher_embeddings

# 正常创建DataLoader即可，此时批量标签会直接是数组切片，转换张量时无性能问题

方案3：修改源码（不推荐） #

直接修改SentenceTransformer.py第547行的代码，将标签列表转为numpy数组后再转张量：

import numpy as np
labels = torch.tensor(np.array(labels))

此方法依赖特定版本的源码，库更新后会失效，不建议采用。

方案1和方案2均可有效解决警告，方案2还能减少训练时的重复计算，进一步提升整体训练效率。

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