PyTorch Dataloader当使用LMDB数据时num_worker＞0引发TypeError: can‘t pickle Environment objects问题的解决

关注

翻译最新推荐文章于 2024-03-10 14:49:57 发布 2.2k 阅读 AI 写同款· GEO 优化 ›

华为昇腾 910B 实战：DeepSeek-R1 蒸馏模型高效部署与性能调优全解析本文详细解析了在华为昇腾910B硬件平台上高效部署DeepSeek-R1蒸馏系列模型（14B/32B版本）的完整实战流程。内容涵盖从MindIE环境搭建、镜像获取、模型准备，到服务配置、内存参数调优、性能测试及常见问题排查，为开发者提供了基于昇腾910B进行大模型部署与优化的全链路指南。阅读详情

PyTorch Dataloader当使用LMDB数据时num_worker>0引发TypeError: can't pickle Environment objects问题的解决

防忘记，原issue：https://github.com/pytorch/vision/issues/689#issuecomment-787215916

解决方法：

不要在__init__方法中调用lmdb.open方法；
在第一次加载数据时打开lmdb。

代码如下

class DataLoader(torch.utils.data.Dataset):
    def __init__(self):
        """do not open lmdb here!!"""

    def open_lmdb(self):
         self.env = lmdb.open(self.lmdb_dir, readonly=True, create=False)
         self.txn = self.env.begin(buffers=True)

    def __getitem__(self, item: int):
        if not hasattr(self, 'txn'):
            self.open_lmdb()
        """
        Then do anything you want with env/txn here.
        """

解释（照搬原文，懒得翻译了）：
The multi-processing actually happens when you create the data iterator (e.g., when calling for datum in dataloader:):
https://github.com/pytorch/pytorch/blob/461014d54b3981c8fa6617f90ff7b7df51ab1e85/torch/utils/data/dataloader.py#L712-L720
In short, it would create multiple processes which “copy” the state of the current process. This copy involves a pickle of the LMDB’s Env thus causes an issue. In our solution, we open it at the first data iteration and the opened lmdb file object would be dedicated to each subprocess.

收录于