python - 如何保存 torchtext 数据集？

Question

我正在处理文本并使用 torchtext.data.Dataset .创建数据集需要花费大量时间。对于只是运行程序，这仍然是可以接受的。但我想调试神经网络的 torch 代码。如果 python 在 Debug模式下启动，数据集创建大约需要 20 分钟 (!!)。这只是为了获得一个工作环境，我可以在其中逐步调试神经网络代码。

我想保存数据集，例如用 pickle 。此示例代码取自 here , 但我删除了此示例不需要的所有内容:

from torchtext import data
from fastai.nlp import *

PATH = 'data/aclImdb/'

TRN_PATH = 'train/all/'
VAL_PATH = 'test/all/'
TRN = f'{PATH}{TRN_PATH}'
VAL = f'{PATH}{VAL_PATH}'

TEXT = data.Field(lower=True, tokenize="spacy")

bs = 64;
bptt = 70

FILES = dict(train=TRN_PATH, validation=VAL_PATH, test=VAL_PATH)
md = LanguageModelData.from_text_files(PATH, TEXT, **FILES, bs=bs, bptt=bptt, min_freq=10)

with open("md.pkl", "wb") as file:
    pickle.dump(md, file)

要运行代码，您需要 aclImdb 数据集，可以从 here 下载.将其提取到此代码段旁边的 data/ 文件夹中。该代码在使用 pickle 的最后一行产生错误:

Traceback (most recent call last):
  File "/home/lhk/programming/fastai_sandbox/lesson4-imdb2.py", line 27, in <module>
    pickle.dump(md, file)
TypeError: 'generator' object is not callable

fastai的样本经常使用dill而不是 pickle 。但这对我也不起作用。

Answer 1

我自己想出了以下函数:

import dill
from pathlib import Path

import torch
from torchtext.data import Dataset

def save_dataset(dataset, path):
    if not isinstance(path, Path):
        path = Path(path)
    path.mkdir(parents=True, exist_ok=True)
    torch.save(dataset.examples, path/"examples.pkl", pickle_module=dill)
    torch.save(dataset.fields, path/"fields.pkl", pickle_module=dill)

def load_dataset(path):
    if not isinstance(path, Path):
        path = Path(path)
    examples = torch.load(path/"examples.pkl", pickle_module=dill)
    fields = torch.load(path/"fields.pkl", pickle_module=dill)
    return Dataset(examples, fields)

并非实际对象可能有点不同，例如，如果您保存 TabularDataset，则 load_dataset 返回类 Dataset 的实例.这不太可能影响数据管道，但可能需要额外努力进行测试。在自定义分词器的情况下，它也应该是可序列化的(例如，没有 lambda 函数等)。