保存PyTorch数据集至磁盘报错：无法序列化_io.BufferedReader对象

错误原因 #

报错TypeError: cannot pickle '_io.BufferedReader' object的核心原因是：你要保存的R3DSemanticDataset对象中包含了无法被序列化的文件句柄（_io.BufferedReader）。torch.save基于pickle实现对象序列化，而文件句柄是绑定当前进程的实时系统资源，pickle无法对其进行序列化存储。

解决方案 #

方案1：保存数据集元数据，重新初始化（推荐） #

不要直接保存整个Dataset对象，而是保存初始化Dataset所需的配置参数，下次使用时重新创建Dataset实例：

# 保存元数据
torch.save({
    'data_root': dataset.data_root,
    'split': dataset.split,
    # 补充其他初始化R3DSemanticDataset需要的参数
}, '../nyu_views_meta.pth')

# 加载时重建Dataset
from clip_fields.datasets import R3DSemanticDataset  # 根据实际模块路径调整
meta = torch.load('../nyu_views_meta.pth')
dataset = R3DSemanticDataset(data_root=meta['data_root'], split=meta['split'])

方案2：自定义Dataset的序列化逻辑 #

修改R3DSemanticDataset类，通过__getstate__和__setstate__方法控制序列化时的对象状态，移除无法序列化的属性，加载时再重新初始化：

import _io
from torch.utils.data import Dataset

class R3DSemanticDataset(Dataset):
    # 原有类代码...

    def __getstate__(self):
        # 复制对象状态，删除无法序列化的IO相关属性
        state = self.__dict__.copy()
        # 遍历删除所有文件句柄类型的属性
        for key in list(state.keys()):
            if isinstance(state[key], _io.BufferedReader):
                del state[key]
        return state

    def __setstate__(self, state):
        # 恢复对象状态，重新初始化之前删除的IO资源
        self.__dict__.update(state)
        # 调用类中初始化文件资源的方法（替换为实际的初始化逻辑）
        self._setup_data_files()

方案3：提取并保存已加载的数据 #

如果只需要保存Dataset中已加载的样本数据，而非Dataset对象本身，可以遍历提取所有样本后保存：

# 提取所有样本数据
dataset_data = [dataset[idx] for idx in range(len(dataset))]

# 保存提取的数据
torch.save(dataset_data, '../nyu_views_data.pth')

# 加载使用
loaded_data = torch.load('../nyu_views_data.pth')
# 若需要Dataset格式，可基于loaded_data封装简单的自定义Dataset类

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