PyTorchDataLoaderLambda序列化问题解决

理解 Can't pickle local object '' 错误

当您在 PyTorch 中使用 DataLoader 并设置 num_workers > 0 时，PyTorch 会启动多个子进程来并行加载数据。为了在主进程和子进程之间传递对象（例如数据集、转换函数等），Python 的 pickle 模块会被用于序列化和反序列化这些对象。然而，pickle 模块对某些类型的对象存在限制，其中一个常见限制就是无法序列化在函数内部定义的本地匿名函数（lambda 函数）或某些复杂的本地闭包。

在提供的错误堆栈跟踪中，我们可以看到问题发生在 _MultiProcessingDataLoaderIter 尝试通过 ForkingPickler 序列化一个对象时，最终抛出了 AttributeError: Can't pickle local object 'get_tokenizer..'。这明确指出，get_tokenizer 函数返回的一个 lambda 对象是导致序列化失败的根本原因。当 DataLoader 尝试将包含此 lambda 函数的数据集或其相关组件传递给子进程时，pickle 无法识别并序列化这个本地定义的匿名函数，从而导致程序崩溃。

分析 get_tokenizer 函数中的问题

提供的 get_tokenizer 函数是一个灵活的工具，用于根据不同的字符串输入返回相应的文本分词器。仔细检查该函数，可以发现以下两个分支会返回 lambda 函数：

1. 当 tokenizer 为 "spacy" 时：

   return lambda s: [tok.text for tok in spacy_en.tokenizer(s)]

   这里返回了一个匿名 lambda 函数，它捕获了 spacy_en 对象的 tokenizer 方法。

2. 当 tokenizer 为 "subword" 时：

   return lambda x: revtok.tokenize(x, decap=True)

   同样，这里也返回了一个匿名 lambda 函数，它封装了 revtok.tokenize 的调用。

这两个 lambda 函数都是在 get_tokenizer 函数被调用时在局部作用域内创建的。当 DataLoader 尝试将这些 lambda 函数（可能通过数据集的 transform 或 collate_fn 间接引用）发送给子进程时，pickle 无法对其进行序列化，从而引发了 AttributeError。

解决方案：将 lambda 替换为命名函数

解决此问题的核心思想是将那些导致序列化失败的本地 lambda 函数替换为具名的函数。具名函数（无论是模块级函数还是嵌套函数）通常能够被 pickle 正确序列化，因为它们具有明确的引用路径和定义。

以下是优化 get_tokenizer 函数的具体步骤，将 lambda 函数替换为嵌套的命名函数：

示例代码：优化 get_tokenizer 函数

import spacy
from nltk.tokenize.moses import MosesTokenizer
import revtok

def get_tokenizer(tokenizer_name):
    """
    根据指定的名称返回一个文本分词器。
    此版本已优化，避免返回不可序列化的本地 lambda 函数。
    """
    if callable(tokenizer_name):
        # 如果传入的已经是可调用对象，则直接返回
        return tokenizer_name

    if tokenizer_name == "spacy":
        try:
            # 导入并加载 SpaCy 模型
            # 注意：在多进程环境下，spacy_en 对象可能会在每个子进程中重新加载，
            # 对于大型模型，这可能导致内存开销。
            spacy_en = spacy.load('en_core_web_sm')
            print("正在加载 SpaCy 分词器模型...")

            # 将 lambda 函数替换为嵌套的命名函数
            def spacy_text_tokenizer(s):
                """使用 SpaCy 模型进行分词的具名函数。"""
                return [tok.text for tok in spacy_en.tokenizer(s)]
            return spacy_text_tokenizer

        except ImportError:
            print("请安装 SpaCy 库和英文分词模型。详情请参考 https://spacy.io")
            raise
        except AttributeError:
            print("请安装 SpaCy 库和英文分词模型。详情请参考 https://spacy.io")
            raise

    elif tokenizer_name == "moses":
        try:
            moses_tokenizer = MosesTokenizer()
            # MosesTokenizer 的 tokenize 方法通常是可序列化的
            return moses_tokenizer.tokenize
        except ImportError:
            print("请安装 NLTK 库。详情请参考 http://nltk.org")
            raise
        except LookupError:
            print("请安装必要的 NLTK 语料库。详情请参考 http://nltk.org")
            raise

    elif tokenizer_name == 'revtok':
        try:
            # revtok.tokenize 是一个模块级函数，通常是可序列化的
            return revtok.tokenize
        except ImportError:
            print("请安装 revtok 库。")
            raise

    elif tokenizer_name == 'subword':
        try:
            # 将 lambda 函数替换为嵌套的命名函数
            def revtok_subword_tokenizer(x):
                """使用 revtok 进行子词分词的具名函数。"""
                return revtok.tokenize(x, decap=True)
            return revtok_subword_tokenizer
        except ImportError:
            print("请安装 revtok 库。")
            raise

    raise ValueError(f"请求的分词器 '{tokenizer_name}' 无效。有效选项包括一个接受字符串的 callable 对象，"
                     "\"revtok\" (用于 revtok 可逆分词器), \"subword\" (用于 revtok 大小写敏感分词器),"
                     "\"spacy\" (用于 SpaCy 英文分词器), 或 \"moses\" (用于 NLTK 的 Moses 分词器)。")

# 示例用法 (假设在你的主脚本中):
# text_field.tokenizer = get_tokenizer(args.tokenizer_type)

修改说明：

• 在 tokenizer_name == "spacy" 分支中，我们定义了一个名为 spacy_text_tokenizer 的嵌套函数来替代原来的 lambda。这个函数捕获了 spacy_en 对象，并执行相同的分词逻辑。
• 在 tokenizer_name == "subword" 分支中，我们同样定义了一个名为 revtok_subword_tokenizer 的嵌套函数来替代 lambda。
• 其他分支（如 "moses" 和 "revtok"）返回的是类实例的方法或模块级函数，这些通常本身就是可序列化的，因此无需修改。

通过这种方式，我们消除了 DataLoader 尝试序列化本地 lambda 函数的根源，从而解决了 AttributeError。

注意事项与最佳实践

在多进程环境中处理可调用对象和资源加载时，还有一些重要的最佳实践需要考虑：

1. 资源初始化与多进程：

   • SpaCy 模型加载： 在上述示例中，spacy.load('en_core_web_sm') 发生在 get_tokenizer 函数内部。如果 DataLoader 的 num_workers > 0，这意味着每个子进程在首次调用 get_tokenizer 时都会加载一次 SpaCy 模型。对于大型模型，这可能导致显著的内存开销和启动延迟。
   • 优化方法：
     • worker_init_fn： PyTorch DataLoader 提供了 worker_init_fn 参数。您可以在这个函数中为每个子进程单独加载和初始化资源（如 SpaCy 模型），并将其存储在进程本地的全局变量中，确保每个工作进程只加载一次。
     • 模块级缓存： 可以设计一个模块级的缓存机制，确保模型只在每个进程中加载一次。

   # 示例 worker_init_fn
   import spacy
   _spacy_model_cache = {}
   
   def worker_init_fn(worker_id):
       global _spacy_model_cache
       if 'en_core_web_sm' not in _spacy_model_cache:
           _spacy_model_cache['en_core_web_sm'] = spacy.load('en_core_web_sm')
       # 可以将 _spacy_model_cache 传递给 dataset 或 transform
       # 例如，通过修改 dataset 的属性，如果 dataset 支持

   然后修改 spacy_text_tokenizer，使其从 _spacy_model_cache 中获取模型。

2. 可序列化性原则：

   • 模块级函数和类： 优先使用模块级定义的函数或类的实例方法作为可调用对象。这些通常具有更好的序列化兼容性。
   • 避免闭包捕获复杂状态： 尽量避免具名函数（即使是嵌套函数）捕获复杂的、不可序列化的外部状态。如果必须捕获，请确保被捕获的对象本身是可序列化的。

3. 调试技巧：num_workers=0：

   • 如果您遇到与多进程相关的序列化错误，一个快速的调试方法是将 DataLoader 的 num_workers 设置为 0。这会强制数据加载在主进程中进行，从而绕过 pickle 机制。如果问题消失，则表明确实是序列化问题。但请记住，这只是一个调试手段，不应作为生产环境的最终解决方案，因为它会牺牲并行加载带来的性能优势。

本篇关于《PyTorchDataLoaderLambda序列化问题解决》的介绍就到此结束啦，但是学无止境，想要了解学习更多关于文章的相关知识，请关注golang学习网公众号！