如何保存DistilBertTokenizer处理后的BatchEncoding以避免Colab会话丢失？

我用Transformers库结合DistilBert做文本分类，数据集有70万行，体量很大，当前在Google Colab上运行代码。构建模型前执行了以下代码：

X = dfreadtrain['review_text'].values
y = dfreadtrain['rating'].values
x_train, x_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=train_y, random_state=42, shuffle=True)
tokenizer = DistilBertTokenizer.from_pretrained(MODEL_NAME)
train_encodings = tokenizer(list(x_train),truncation=True,padding=True)
test_encodings = tokenizer(list(x_test), truncation=True, padding=True)
print(type(train_encodings))

这段代码执行耗时极长，但Colab会话一旦中断，生成的train_encodings和test_encodings（类型为transformers.tokenization_utils_base.BatchEncoding）就会丢失，请问能不能把这两个编码对象保存到文件中？

解决方案 #

完全可以，以下是两种常用的保存与加载方法：

方法一：使用BatchEncoding自带的save_to_disk方法（推荐） #

Transformers库的BatchEncoding对象内置了官方推荐的save_to_disk方法，可直接将编码数据保存到指定目录，加载时通过tokenizer.load_from_disk恢复：

保存代码

# 将编码对象保存到本地目录
train_encodings.save_to_disk("./train_encodings")
test_encodings.save_to_disk("./test_encodings")

加载代码

# 从磁盘加载编码对象
train_encodings = tokenizer.load_from_disk("./train_encodings")
test_encodings = tokenizer.load_from_disk("./test_encodings")

该方法保存的结构清晰，兼容性强，尤其适合大规模数据场景。

方法二：使用pickle序列化保存 #

如果习惯Python通用的序列化方式，可通过pickle将BatchEncoding对象保存为二进制文件：

保存代码

import pickle

# 保存为pickle格式文件
with open("train_encodings.pkl", "wb") as f:
    pickle.dump(train_encodings, f)

with open("test_encodings.pkl", "wb") as f:
    pickle.dump(test_encodings, f)

加载代码

import pickle

# 从pickle文件加载编码对象
with open("train_encodings.pkl", "rb") as f:
    train_encodings = pickle.load(f)

with open("test_encodings.pkl", "rb") as f:
    test_encodings = pickle.load(f)

注意：这种方法存在版本兼容性风险（不同Python或Transformers版本可能无法正常加载），建议优先使用方法一。

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