基于aladdinpersson/Machine-Learning-Collection项目的PyTorch文本数据集处理教程-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00895/article/details/148441313

基于aladdinpersson/Machine-Learning-Collection项目的PyTorch文本数据集处理教程

Machine-Learning-Collection A resource for learning about Machine learning & Deep Learning 项目地址: https://gitcode.com/gh_mirrors/ma/Machine-Learning-Collection

在机器学习项目中，处理文本数据是一个常见但具有挑战性的任务。本文将深入解析如何使用PyTorch构建自定义文本数据集处理流程，特别适合初学者理解文本数据处理的基本原理。

文本数据处理的核心挑战

处理文本数据时，我们需要解决几个关键问题：

如何将自然语言转换为数值表示（词向量）
如何构建词汇表（Vocabulary）
如何处理变长文本序列
如何与图像数据等其他模态数据结合处理

词汇表(Vocabulary)实现

词汇表是文本处理的基础组件，负责单词与索引之间的双向映射：

class Vocabulary:
    def __init__(self, freq_threshold):
        self.itos = {0: "<PAD>", 1: "<SOS>", 2: "<EOS>", 3: "<UNK>"}
        self.stoi = {"<PAD>": 0, "<SOS>": 1, "<EOS>": 2, "<UNK>": 3}
        self.freq_threshold = freq_threshold

itos (index to string): 索引到单词的映射
stoi (string to index): 单词到索引的映射
特殊标记说明：
- <PAD>: 填充标记，用于统一序列长度
- <SOS>: 句子开始标记
- <EOS>: 句子结束标记
- <UNK>: 未知单词标记

词汇表构建时采用频率阈值(freq_threshold)来控制词汇量，只有当单词出现次数达到阈值才会被加入词汇表。

文本分词处理

使用spacy库进行英语分词：

@staticmethod
def tokenizer_eng(text):
    return [tok.text.lower() for tok in spacy_eng.tokenizer(text)]

这种分词方式考虑了英语的语言特性，能够正确处理缩写、连字符等情况。

数据集类实现

FlickrDataset类继承自PyTorch的Dataset，用于处理图像-文本对数据：

class FlickrDataset(Dataset):
    def __init__(self, root_dir, captions_file, transform=None, freq_threshold=5):
        self.root_dir = root_dir
        self.df = pd.read_csv(captions_file)
        self.transform = transform
        self.imgs = self.df["image"]
        self.captions = self.df["caption"]
        self.vocab = Vocabulary(freq_threshold)
        self.vocab.build_vocabulary(self.captions.tolist())

关键点：

同时加载图像和文本数据
初始化词汇表并基于所有文本构建词汇
支持图像变换(transform)

批处理与填充

文本数据的变长特性要求我们实现自定义的批处理逻辑：

class MyCollate:
    def __init__(self, pad_idx):
        self.pad_idx = pad_idx

    def __call__(self, batch):
        imgs = [item[0].unsqueeze(0) for item in batch]
        imgs = torch.cat(imgs, dim=0)
        targets = [item[1] for item in batch]
        targets = pad_sequence(targets, batch_first=False, padding_value=self.pad_idx)
        return imgs, targets

使用pad_sequence函数对变长序列进行填充，使一个batch内的所有序列具有相同长度。

数据加载器配置

get_loader函数封装了数据加载的完整流程：

def get_loader(
    root_folder,
    annotation_file,
    transform,
    batch_size=32,
    num_workers=8,
    shuffle=True,
    pin_memory=True,
):
    dataset = FlickrDataset(root_folder, annotation_file, transform=transform)
    pad_idx = dataset.vocab.stoi["<PAD>"]
    loader = DataLoader(
        dataset=dataset,
        batch_size=batch_size,
        num_workers=num_workers,
        shuffle=shuffle,
        pin_memory=pin_memory,
        collate_fn=MyCollate(pad_idx=pad_idx),
    )
    return loader, dataset

参数说明：

num_workers: 多线程加载数据
pin_memory: 将数据固定在内存中，加速GPU传输
collate_fn: 指定自定义的批处理函数

实际应用示例

if __name__ == "__main__":
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
    ])

    loader, dataset = get_loader(
        "flickr8k/images/", "flickr8k/captions.txt", transform=transform
    )

    for idx, (imgs, captions) in enumerate(loader):
        print(imgs.shape)  # 打印图像张量形状
        print(captions.shape)  # 打印文本张量形状

这个示例展示了如何将图像调整为224x224大小并转换为张量，同时加载对应的文本描述。