评价此页

torch.package#

创建日期:2025年6月10日 | 最后更新日期:2025年7月15日

torch.package 增加了对创建包含工件(artifacts)和任意 PyTorch 代码的程序包的支持。这些程序包可以被保存、共享、用于在以后或不同机器上加载和执行模型,甚至可以使用 torch::deploy 部署到生产环境。

本文档包含教程、操作指南、解释和 API 参考,将帮助您了解更多关于 torch.package 及其使用方法的信息。

警告

此模块依赖于不安全的 pickle 模块。请仅解包您信任的数据。

构建恶意的 pickle 数据是可能的,这会在反序列化(unpickling)过程中执行任意代码。切勿解包来自不受信任来源的数据,或可能已被篡改的数据。

有关详细信息,请查阅 pickle 模块的官方文档

教程#

打包您的第一个模型#

一份指导您打包和解包简单模型的教程可在 Colab 上找到。完成此练习后,您将熟悉创建和使用 Torch 包的基本 API。

如何……#

查看包中有什么内容?#

将包视为 ZIP 归档文件#

torch.package 的容器格式是 ZIP,因此任何适用于标准 ZIP 文件的工具都可用于探索其内容。一些与 ZIP 文件交互的常用方法:

  • unzip my_package.pt 会将 torch.package 归档文件解压到磁盘,您可以在那里自由检查其内容。

$ unzip my_package.pt && tree my_package
my_package
├── .data
│   ├── 94304870911616.storage
│   ├── 94304900784016.storage
│   ├── extern_modules
│   └── version
├── models
│   └── model_1.pkl
└── torchvision
    └── models
        ├── resnet.py
        └── utils.py
~ cd my_package && cat torchvision/models/resnet.py
...
  • Python 的 zipfile 模块提供了一种读取和写入 ZIP 归档内容的标准方法。

from zipfile import ZipFile
with ZipFile("my_package.pt") as myzip:
    file_bytes = myzip.read("torchvision/models/resnet.py")
    # edit file_bytes in some way
    myzip.writestr("torchvision/models/resnet.py", new_file_bytes)
  • vim 具有原生读取 ZIP 归档的能力。您甚至可以编辑文件并将它们 :write 回归档中!

# add this to your .vimrc to treat `*.pt` files as zip files
au BufReadCmd *.pt call zip#Browse(expand("<amatch>"))

~ vi my_package.pt

使用 file_structure() API#

PackageImporter 提供了一个 file_structure() 方法,该方法将返回一个可打印且可查询的 Directory 对象。Directory 对象是一个简单的目录结构,您可以利用它来探索 torch.package 的当前内容。

Directory 对象本身可直接打印,并会打印出文件树表示。要过滤返回的内容,请使用类 glob 的 includeexclude 过滤参数。

with PackageExporter('my_package.pt') as pe:
    pe.save_pickle('models', 'model_1.pkl', mod)

importer = PackageImporter('my_package.pt')
# can limit printed items with include/exclude args
print(importer.file_structure(include=["**/utils.py", "**/*.pkl"], exclude="**/*.storage"))
print(importer.file_structure()) # will print out all files

输出

# filtered with glob pattern:
#    include=["**/utils.py", "**/*.pkl"], exclude="**/*.storage"
─── my_package.pt
    ├── models
    │   └── model_1.pkl
    └── torchvision
        └── models
            └── utils.py

# all files
─── my_package.pt
    ├── .data
    │   ├── 94304870911616.storage
    │   ├── 94304900784016.storage
    │   ├── extern_modules
    │   └── version
    ├── models
    │   └── model_1.pkl
    └── torchvision
        └── models
            ├── resnet.py
            └── utils.py

您还可以使用 has_file() 方法查询 Directory 对象。

importer_file_structure = importer.file_structure()
found: bool = importer_file_structure.has_file("package_a/subpackage.py")

查看为何某个给定模块被作为依赖项包含在内?#

假设存在一个给定模块 foo,并且您想知道为什么您的 PackageExporterfoo 作为依赖项引入。

PackageExporter.get_rdeps() 将返回所有直接依赖于 foo 的模块。

如果您想查看给定模块 src 是如何依赖于 foo 的,PackageExporter.all_paths() 方法将返回一个 DOT 格式的图,显示 srcfoo 之间的所有依赖路径。

如果您只是想查看 :class:PackageExporter 的完整依赖图,可以使用 PackageExporter.dependency_graph_string()

将任意资源包含在我的包中并在稍后访问它们?#

PackageExporter 公开了三个方法:save_picklesave_textsave_binary,允许您将 Python 对象、文本和二进制数据保存到包中。

with torch.PackageExporter("package.pt") as exporter:
    # Pickles the object and saves to `my_resources/tensor.pkl` in the archive.
    exporter.save_pickle("my_resources", "tensor.pkl", torch.randn(4))
    exporter.save_text("config_stuff", "words.txt", "a sample string")
    exporter.save_binary("raw_data", "binary", my_bytes)

PackageImporter 公开了名为 load_pickleload_textload_binary 的互补方法,允许您从包中加载 Python 对象、文本和二进制数据。

importer = torch.PackageImporter("package.pt")
my_tensor = importer.load_pickle("my_resources", "tensor.pkl")
text = importer.load_text("config_stuff", "words.txt")
binary = importer.load_binary("raw_data", "binary")

自定义类的打包方式?#

torch.package 允许自定义类的打包方式。此行为通过在类上定义 __reduce_package__ 方法以及定义相应的解包函数来访问。这类似于为 Python 的正常序列化过程定义 __reduce__

步骤

  1. 在目标类上定义 __reduce_package__(self, exporter: PackageExporter) 方法。此方法应执行将类实例保存到包中的工作,并返回一个元组,其中包含调用解包函数所需的参数。当 PackageExporter 遇到目标类的实例时,会调用此方法。

  2. 为该类定义一个解包函数。此解包函数应负责重构并返回该类的实例。函数签名的第一个参数应该是 PackageImporter 实例,其余参数由用户定义。

# foo.py [Example of customizing how class Foo is packaged]
from torch.package import PackageExporter, PackageImporter
import time


class Foo:
    def __init__(self, my_string: str):
        super().__init__()
        self.my_string = my_string
        self.time_imported = 0
        self.time_exported = 0

    def __reduce_package__(self, exporter: PackageExporter):
        """
        Called by ``torch.package.PackageExporter``'s Pickler's ``persistent_id`` when
        saving an instance of this object. This method should do the work to save this
        object inside of the ``torch.package`` archive.

        Returns function w/ arguments to load the object from a
        ``torch.package.PackageImporter``'s Pickler's ``persistent_load`` function.
        """

        # use this pattern to ensure no naming conflicts with normal dependencies,
        # anything saved under this module name shouldn't conflict with other
        # items in the package
        generated_module_name = f"foo-generated._{exporter.get_unique_id()}"
        exporter.save_text(
            generated_module_name,
            "foo.txt",
            self.my_string + ", with exporter modification!",
        )
        time_exported = time.clock_gettime(1)

        # returns de-packaging function w/ arguments to invoke with
        return (unpackage_foo, (generated_module_name, time_exported,))


def unpackage_foo(
    importer: PackageImporter, generated_module_name: str, time_exported: float
) -> Foo:
    """
    Called by ``torch.package.PackageImporter``'s Pickler's ``persistent_load`` function
    when depickling a Foo object.
    Performs work of loading and returning a Foo instance from a ``torch.package`` archive.
    """
    time_imported = time.clock_gettime(1)
    foo = Foo(importer.load_text(generated_module_name, "foo.txt"))
    foo.time_imported = time_imported
    foo.time_exported = time_exported
    return foo

# example of saving instances of class Foo

import torch
from torch.package import PackageImporter, PackageExporter
import foo

foo_1 = foo.Foo("foo_1 initial string")
foo_2 = foo.Foo("foo_2 initial string")
with PackageExporter('foo_package.pt') as pe:
    # save as normal, no extra work necessary
    pe.save_pickle('foo_collection', 'foo1.pkl', foo_1)
    pe.save_pickle('foo_collection', 'foo2.pkl', foo_2)

pi = PackageImporter('foo_package.pt')
print(pi.file_structure())
imported_foo = pi.load_pickle('foo_collection', 'foo1.pkl')
print(f"foo_1 string: '{imported_foo.my_string}'")
print(f"foo_1 export time: {imported_foo.time_exported}")
print(f"foo_1 import time: {imported_foo.time_imported}")
# output of running above script
─── foo_package
    ├── foo-generated
    │   ├── _0
    │   │   └── foo.txt
    │   └── _1
    │       └── foo.txt
    ├── foo_collection
    │   ├── foo1.pkl
    │   └── foo2.pkl
    └── foo.py

foo_1 string: 'foo_1 initial string, with reduction modification!'
foo_1 export time: 9857706.650140837
foo_1 import time: 9857706.652698385

在源代码中测试它是否在包内运行?#

PackageImporter 会将其初始化的每个模块添加 __torch_package__ 属性。您的代码可以检查此属性是否存在,以确定它是在已打包的上下文中运行还是在常规上下文中运行。

# In foo/bar.py:

if "__torch_package__" in dir():  # true if the code is being loaded from a package
    def is_in_package():
        return True

    UserException = Exception
else:
    def is_in_package():
        return False

    UserException = UnpackageableException

现在,代码将根据它是通过您的 Python 环境正常导入还是从 torch.package 导入而表现出不同的行为。

from foo.bar import is_in_package

print(is_in_package())  # False

loaded_module = PackageImporter(my_package).import_module("foo.bar")
loaded_module.is_in_package()  # True

警告:通常情况下,让代码根据是否打包而表现不同是不好的做法。这可能导致难以调试的问题,这些问题对您如何导入代码非常敏感。如果您的包旨在被广泛使用,请考虑重构代码,使其无论如何加载都表现相同。

将补丁代码植入包中?#

PackageExporter 提供了一个 save_source_string() 方法,允许用户将任意 Python 源代码保存到您选择的模块中。

with PackageExporter(f) as exporter:
    # Save the my_module.foo available in your current Python environment.
    exporter.save_module("my_module.foo")

    # This saves the provided string to my_module/foo.py in the package archive.
    # It will override the my_module.foo that was previously saved.
    exporter.save_source_string("my_module.foo", textwrap.dedent(
        """\
        def my_function():
            print('hello world')
        """
    ))

    # If you want to treat my_module.bar as a package
    # (e.g. save to `my_module/bar/__init__.py` instead of `my_module/bar.py)
    # pass is_package=True,
    exporter.save_source_string("my_module.bar",
                                "def foo(): print('hello')\n",
                                is_package=True)

importer = PackageImporter(f)
importer.import_module("my_module.foo").my_function()  # prints 'hello world'

从已打包的代码中访问包内容?#

PackageImporter 实现了 importlib.resources API,用于从包内访问资源。

with PackageExporter(f) as exporter:
    # saves text to my_resource/a.txt in the archive
    exporter.save_text("my_resource", "a.txt", "hello world!")
    # saves the tensor to my_pickle/obj.pkl
    exporter.save_pickle("my_pickle", "obj.pkl", torch.ones(2, 2))

    # see below for module contents
    exporter.save_module("foo")
    exporter.save_module("bar")

importlib.resources API 允许从已打包的代码内访问资源。

# foo.py:
import importlib.resources
import my_resource

# returns "hello world!"
def get_my_resource():
    return importlib.resources.read_text(my_resource, "a.txt")

使用 importlib.resources 是从已打包代码中访问包内容的推荐方式,因为它符合 Python 标准。但是,也可以在已打包的代码中访问父级 :class:PackageImporter 实例本身。

# bar.py:
import torch_package_importer # this is the PackageImporter that imported this module.

# Prints "hello world!", equivalent to importlib.resources.read_text
def get_my_resource():
    return torch_package_importer.load_text("my_resource", "a.txt")

# You also do things that the importlib.resources API does not support, like loading
# a pickled object from the package.
def get_my_pickle():
    return torch_package_importer.load_pickle("my_pickle", "obj.pkl")

区分已打包代码和非打包代码?#

要判断对象的代码是否来自 torch.package,请使用 torch.package.is_from_package() 函数。注意:如果对象来自某个包,但其定义来自标记为 extern 的模块或来自 stdlib,则此检查将返回 False

importer = PackageImporter(f)
mod = importer.import_module('foo')
obj = importer.load_pickle('model', 'model.pkl')
txt = importer.load_text('text', 'my_test.txt')

assert is_from_package(mod)
assert is_from_package(obj)
assert not is_from_package(txt) # str is from stdlib, so this will return False

重新导出已导入的对象?#

要重新导出之前由 PackageImporter 导入的对象,您必须使新的 PackageExporter 知道原始的 PackageImporter,以便它能够找到您对象依赖项的源代码。

importer = PackageImporter(f)
obj = importer.load_pickle("model", "model.pkl")

# re-export obj in a new package
with PackageExporter(f2, importer=(importer, sys_importer)) as exporter:
    exporter.save_pickle("model", "model.pkl", obj)

解释#

torch.package 格式概述#

torch.package 文件是一个 ZIP 归档文件,通常使用 .pt 扩展名。在 ZIP 归档内部,有两种文件:

  • 框架文件,放置在 .data/ 中。

  • 用户文件,即其他所有内容。

例如,这是一个来自 torchvision 的完全打包的 ResNet 模型的样子:

resnet
├── .data  # All framework-specific data is stored here.
│   │      # It's named to avoid conflicts with user-serialized code.
│   ├── 94286146172688.storage  # tensor data
│   ├── 94286146172784.storage
│   ├── extern_modules  # text file with names of extern modules (e.g. 'torch')
│   ├── version         # version metadata
│   ├── ...
├── model  # the pickled model
│   └── model.pkl
└── torchvision  # all code dependencies are captured as source files
    └── models
        ├── resnet.py
        └── utils.py

框架文件#

.data/ 目录由 torch.package 拥有,其内容被视为私有实现细节。torch.package 格式不对 .data/ 的内容做任何保证,但所做的任何更改都将向后兼容(即较新版本的 PyTorch 始终能够加载较旧的 torch.package)。

目前,.data/ 目录包含以下项目:

  • version:序列化格式的版本号,以便 torch.package 导入基础设施知道如何加载此包。

  • extern_modules:被视为 extern 的模块列表。extern 模块将使用加载环境的系统导入器进行导入。

  • *.storage:序列化的张量数据。

.data
├── 94286146172688.storage
├── 94286146172784.storage
├── extern_modules
├── version
├── ...

用户文件#

归档中的所有其他文件均由用户放入。其布局与 Python 常规包 完全相同。有关 Python 打包工作原理的深入研究,请参阅 此文章(它稍微有点过时,请对照 Python 参考文档 核对实现细节)。

<package root>
├── model  # the pickled model
│   └── model.pkl
├── another_package
│   ├── __init__.py
│   ├── foo.txt         # a resource file , see importlib.resources
│   └── ...
└── torchvision
    └── models
        ├── resnet.py   # torchvision.models.resnet
        └── utils.py    # torchvision.models.utils

torch.package 如何查找您代码的依赖项#

分析对象的依赖项#

当您发出 save_pickle(obj, ...) 调用时,PackageExporter 将正常序列化对象。然后,它使用 pickletools 标准库模块解析 pickle 字节码。

在 pickle 中,对象与描述如何查找对象类型实现的 GLOBAL 操作码一起保存,例如:

GLOBAL 'torchvision.models.resnet Resnet`

依赖解析器将收集所有 GLOBAL 操作并将其标记为已序列化对象的依赖项。有关序列化和 pickle 格式的更多信息,请查阅 Python 文档

分析模块的依赖项#

当一个 Python 模块被识别为依赖项时,torch.package 会遍历模块的 Python AST 表示,并查找完全支持标准形式的 import 语句:from x import yimport zfrom w import v as u 等。当遇到这些 import 语句时,torch.package 会将导入的模块注册为依赖项,然后这些依赖项也会以相同的 AST 遍历方式进行解析。

注意:AST 解析对 __import__(...) 语法的支持有限,且不支持 importlib.import_module 调用。通常,您不应期望 torch.package 能检测到动态导入。

依赖管理#

torch.package 会自动查找您的代码和对象所依赖的 Python 模块。此过程称为依赖解析。对于依赖解析器找到的每个模块,您必须指定一个要采取的操作

允许的操作包括:

  • intern:将此模块放入包中。

  • extern:将此模块声明为包的外部依赖项。

  • mock:存根(stub)化此模块。

  • deny:依赖此模块将在包导出过程中引发错误。

最后,还有一个重要的操作,在技术上不属于 torch.package

  • 重构:删除或更改代码中的依赖项。

请注意,操作仅在整个 Python 模块上定义。无法“仅”打包模块中的一个函数或类而将其余部分排除在外。这是有意为之的。Python 不提供模块中定义的对象之间的清晰边界。定义的唯一依赖组织单元是模块,因此这就是 torch.package 所使用的单位。

使用模式将操作应用于模块。模式可以是模块名称("foo.bar")或 glob(如 "foo.**")。您可以使用 PackageExporter 上的方法将模式与操作关联,例如:

my_exporter.intern("torchvision.**")
my_exporter.extern("numpy")

如果模块匹配某个模式,则对其应用相应的操作。对于给定模块,将按定义顺序检查模式,并采用第一个匹配的操作。

intern#

如果模块被 intern,它将被放入包中。

此操作适用于您的模型代码,或任何您想要打包的相关代码。例如,如果您尝试从 torchvision 打包一个 ResNet,则需要 intern torchvision.models.resnet 模块。

在包导入时,当已打包的代码尝试导入一个 intern 过的模块时,PackageImporter 将在您的包中查找该模块。如果找不到,将引发错误。这确保了每个 PackageImporter 都与加载环境隔离——即使您的包和加载环境中都存在 my_interned_modulePackageImporter 也只会使用包中的版本。

注意:只有 Python 源模块可以被 intern。其他类型的模块(如 C 扩展模块和字节码模块)如果尝试 intern,将引发错误。这些类型的模块需要进行 mockextern

extern#

如果模块被 extern,它将不会被打包。相反,它将被添加到此包的外部依赖项列表中。您可以在 package_exporter.extern_modules 上找到此列表。

在包导入时,当已打包的代码尝试导入一个 extern 过的模块时,PackageImporter 将使用默认的 Python 导入器查找该模块,就像执行了 importlib.import_module("my_externed_module") 一样。如果找不到,将引发错误。

通过这种方式,您可以从包内依赖第三方库(如 numpyscipy),而无需将它们也打包进去。

警告:如果任何外部库以不向后兼容的方式发生更改,您的包可能会加载失败。如果您需要包的长期可重复性,请尽量限制对 extern 的使用。

mock#

如果模块被 mock,它将不会被打包。取而代之的是,一个存根模块将被打包。该存根模块允许您从中检索对象(因此 from my_mocked_module import foo 不会报错),但任何使用该对象的行为都会引发 NotImplementedError

mock 应仅用于您“确定”在已加载包中不需要,但在非打包内容中仍然可用的代码。例如:初始化/配置代码,或仅用于调试/训练的代码。

警告:通常情况下,mock 应作为最后手段使用。它会导致已打包代码和非打包代码之间的行为差异,这可能会导致后续的困惑。建议优先重构您的代码以删除不需要的依赖项。

重构#

管理依赖的最佳方式是根本没有依赖!通常,代码可以通过重构来消除不必要的依赖项。以下是编写具有清晰依赖关系的规范代码的一些指南(这些通常也是良好的编程实践!):

仅包含您使用的内容。不要在代码中留下未使用的导入。依赖解析器不够聪明,无法判断它们确实未使用,并且会尝试处理它们。

限定您的导入。例如,不要写 import foo 并在稍后使用 foo.bar.baz,更推荐写 from foo.bar import baz。这更精确地指定了您的实际依赖项(foo.bar),并让依赖解析器知道您不需要整个 foo

将具有不相关功能的大文件拆分成较小的文件。如果您的 utils 模块包含大量不相关功能的混合体,那么任何依赖于 utils 的模块都需要引入大量不相关的依赖项,即使您只用到了其中的一小部分。更推荐定义可独立打包的单用途模块。

模式#

模式允许您使用方便的语法指定模块组。模式的语法和行为遵循 Bazel/Buck 的 glob()

我们试图匹配模式的模块称为候选对象(candidate)。候选对象由一系列通过分隔符字符串隔开的片段组成,例如 foo.bar.baz

模式包含一个或多个片段。片段可以是

  • 字面量字符串(例如 foo),进行精确匹配。

  • 包含通配符的字符串(例如 torchfoo*baz*)。通配符匹配任何字符串,包括空字符串。

  • 双通配符(**)。它匹配零个或多个完整片段。

示例

  • torch.**:匹配 torch 及其所有子模块,例如 torch.nntorch.nn.functional

  • torch.*:匹配 torch.nntorch.functional,但不匹配 torch.nn.functionaltorch

  • torch*.**:匹配 torchtorchvision 及其所有子模块

指定操作时,可以传入多个模式,例如

exporter.intern(["torchvision.models.**", "torchvision.utils.**"])

如果模块匹配其中任何一个模式,则该模块将匹配此操作。

还可以指定排除模式,例如

exporter.mock("**", exclude=["torchvision.**"])

如果模块匹配任何排除模式,则不会匹配此操作。在此示例中,我们将模拟除 torchvision 及其子模块之外的所有模块。

当模块可能匹配多个操作时,将采用最先定义的操作。

torch.package 的难点#

避免模块中的全局状态#

Python 可以非常容易地在模块级作用域内绑定对象并运行代码。这通常没问题——毕竟函数和类都是以这种方式绑定到名称的。但是,当您为了修改而定义模块级对象时,情况就会变得复杂,从而引入了可变的全局状态。

可变全局状态非常有用——它可以减少样板代码、允许开放式注册到表中等。但除非非常谨慎地使用,否则在使用 torch.package 时可能会引起复杂的问题。

每个 PackageImporter 都会为其内容创建一个独立的环境。这很好,因为这意味着我们可以加载多个包并确保它们彼此隔离,但当模块以假定共享可变全局状态的方式编写时,这种行为可能会产生难以调试的错误。

类型在包和加载环境之间是不共享的#

您从 PackageImporter 导入的任何类都将是该导入器特有的类版本。例如

from foo import MyClass

my_class_instance = MyClass()

with PackageExporter(f) as exporter:
    exporter.save_module("foo")

importer = PackageImporter(f)
imported_MyClass = importer.import_module("foo").MyClass

assert isinstance(my_class_instance, MyClass)  # works
assert isinstance(my_class_instance, imported_MyClass)  # ERROR!

在此示例中,MyClassimported_MyClass不同的类型。在这个特定示例中,MyClassimported_MyClass 具有完全相同的实现,因此您可能认为将它们视为同一个类是可以的。但考虑一下 imported_MyClass 来自一个带有完全不同 MyClass 实现的旧包的情况——在这种情况下,将它们视为同一个类是不安全的。

在底层,每个导入器都有一个前缀,允许它唯一地标识类

print(MyClass.__name__)  # prints "foo.MyClass"
print(imported_MyClass.__name__)  # prints <torch_package_0>.foo.MyClass

这意味着当其中一个参数来自包而另一个不是时,您不应指望 isinstance 检查能正常工作。如果您需要此功能,请考虑以下选项

  • 使用鸭子类型(直接使用类,而不是显式检查它是否为给定类型)。

  • 使类型关系成为类契约的显式部分。例如,您可以添加一个属性标签 self.handler = "handle_me_this_way",并让客户端代码检查 handler 的值,而不是直接检查类型。

torch.package 如何保持包彼此隔离#

每个 PackageImporter 实例都为其模块和对象创建一个独立的、隔离的环境。包中的模块只能导入其他打包的模块,或标记为 extern 的模块。如果您使用多个 PackageImporter 实例来加载同一个包,您将得到多个互不干扰的独立环境。

这是通过使用自定义导入器扩展 Python 的导入基础结构来实现的。PackageImporter 提供了与 importlib 导入器相同的核心 API;即,它实现了 import_module__import__ 方法。

当您调用 PackageImporter.import_module() 时,PackageImporter 将像系统导入器一样构建并返回一个新模块。然而,PackageImporter 会对返回的模块进行修补,使其使用 self(即该 PackageImporter 实例)通过在包内查找而不是搜索用户的 Python 环境来满足未来的导入请求。

名称混淆(Mangling)#

为了避免混淆(“这个 foo.bar 对象是来自我的包的对象,还是来自我的 Python 环境的对象?”),PackageImporter 通过添加混淆前缀来混淆所有已导入模块的 __name____file__

对于 __name__,名称如 torchvision.models.resnet18 将变为 <torch_package_0>.torchvision.models.resnet18

对于 __file__,名称如 torchvision/models/resnet18.py 将变为 <torch_package_0>.torchvision/modules/resnet18.py

名称混淆有助于避免不同包之间模块名称的无意重名,并通过使堆栈跟踪和打印语句更清晰地显示它们是否引用已打包的代码,从而帮助您进行调试。有关面向开发人员的混淆详情,请查阅 torch/package/ 中的 mangling.md

API 参考#

class torch.package.PackagingError(dependency_graph, debug=False)[source]#

当导出包出现问题时,会引发此异常。PackageExporter 将尝试收集所有错误并一次性呈现给您。

class torch.package.EmptyMatchError[source]#

这是一个异常,当 mock 或 extern 被标记为 allow_empty=False 且在打包过程中未与任何模块匹配时抛出。

class torch.package.PackageExporter(f, importer=<torch.package.importer._SysImporter object>, debug=False)[source]#

导出器允许您将代码包、序列化的 Python 数据以及任意二进制和文本资源写入自包含的包中。

导入器可以以封闭的方式加载此代码,使得代码从包中加载,而不是使用普通的 Python 导入系统。这允许对 PyTorch 模型代码和数据进行打包,以便它们可以在服务器上运行或用于以后的迁移学习。

包中包含的代码在创建时会从原始源逐文件复制,文件格式为特别组织的 zip 文件。包的未来用户可以解压缩包,并编辑代码以对其执行自定义修改。

包的导入器确保模块中的代码只能从包内加载,除非使用 extern() 显式列为外部模块。zip 归档文件中的 extern_modules 文件列出了包外部依赖的所有模块。这可以防止“隐式”依赖——即包在本地运行时因为导入了本地安装的包而正常工作,但在复制到另一台机器时却失败的情况。

向包中添加源代码时,导出器可以选择对其进行扫描以查找进一步的代码依赖项(dependencies=True)。它会查找 import 语句,解析对限定模块名称的相对引用,并执行用户指定的操作(参见:extern()mock()intern())。

__init__(f, importer=<torch.package.importer._SysImporter object>, debug=False)[source]#

创建导出器。

参数:
  • f (str | PathLike[str] | IO[bytes]) – 要导出的位置。可以是包含文件名的 string/Path 对象,或者是二进制 I/O 对象。

  • importer (Importer | Sequence[Importer]) – 如果传入单个 Importer,则使用它搜索模块。如果传入导入器序列,将从中构建一个 OrderedImporter

  • debug (bool) – 如果设置为 True,则将损坏模块的路径添加到 PackagingErrors 中。

add_dependency(module_name, dependencies=True)[source]#

给定一个模块,根据用户指定的模式将其添加到依赖关系图中。

all_paths(src, dst)[source]#
返回子图的点(dot)表示,

其中包含从 src 到 dst 的所有路径。

返回:

包含从 src 到 dst 的所有路径的点(dot)表示。(https://graphviz.cn/doc/info/lang.html)

返回类型:

str

close()[source]#

将包写入文件系统。close() 之后的所有调用现在都是无效的。最好使用资源保护语法(resource guard syntax)代替

with PackageExporter("file.zip") as e:
    ...
denied_modules()[source]#

返回当前被拒绝的所有模块。

返回:

包含在此包中将被拒绝的模块名称的列表。

返回类型:

list[str]

deny(include, *, exclude=())[source]#

将名称匹配给定 glob 模式的模块列入黑名单,使其无法被包导入。如果发现对任何匹配包的依赖,将引发 PackagingError

参数:
  • include (Union[List[str], str]) – 一个字符串,例如 "my_package.my_subpackage",或者要 extern 的模块名称列表。这也可以是 glob 风格的模式,如 mock() 中所述。

  • exclude (Union[List[str], str]) – 一个可选模式,用于排除一些匹配 include 字符串的模式。

dependency_graph_string()[source]#

返回包中依赖关系的有向图字符串表示。

返回:

包中依赖关系的字符串表示。

返回类型:

str

extern(include, *, exclude=(), allow_empty=True)[source]#

module 包含在包可以导入的外部模块列表中。这将防止依赖发现将其保存到包中。导入器将直接从标准导入系统加载外部模块。外部模块的代码也必须存在于加载包的进程中。

参数:
  • include (Union[List[str], str]) – 一个字符串,例如 "my_package.my_subpackage",或者要 extern 的模块名称列表。这也可以是 glob 风格的模式,如 mock() 中所述。

  • exclude (Union[List[str], str]) – 一个可选模式,用于排除一些匹配 include 字符串的模式。

  • allow_empty (bool) – 一个可选标志,用于指定此 extern 方法调用指定的外部模块在打包期间是否必须与某些模块匹配。如果添加了 allow_empty=False 的外部模块 glob 模式,并且在任何模块匹配该模式之前调用了 close()(无论是显式调用还是通过 __exit__),则会抛出异常。如果 allow_empty=True,则不会抛出此类异常。

externed_modules()[source]#

返回当前 extern 的所有模块。

返回:

包含在此包中将被 extern 的模块名称的列表。

返回类型:

list[str]

get_rdeps(module_name)[source]#

返回所有依赖于 module_name 的模块列表。

返回:

包含依赖于 module_name 的模块名称列表。

返回类型:

list[str]

get_unique_id()[source]#

获取一个 ID。此 ID 保证在此包中仅会被分发一次。

返回类型:

str

intern(include, *, exclude=(), allow_empty=True)[source]#

指定应被打包的模块。模块必须匹配某个 intern 模式,才能被包含在包中并递归处理其依赖项。

参数:
  • include (Union[List[str], str]) – 一个字符串,例如 “my_package.my_subpackage”,或者要 extern 的模块名称列表。这也可以是 glob 风格的模式,如 mock() 中所述。

  • exclude (Union[List[str], str]) – 一个可选模式,用于排除一些匹配 include 字符串的模式。

  • allow_empty (bool) – 一个可选标志,用于指定此 intern 方法调用指定的内部模块在打包期间是否必须与某些模块匹配。如果添加了 allow_empty=Falseintern 模块 glob 模式,并且在任何模块匹配该模式之前调用了 close()(无论是显式调用还是通过 __exit__),则会抛出异常。如果 allow_empty=True,则不会抛出此类异常。

interned_modules()[source]#

返回当前 intern 的所有模块。

返回:

包含在此包中将被 intern 的模块名称的列表。

返回类型:

list[str]

mock(include, *, exclude=(), allow_empty=True)[source]#

用模拟实现替换某些必需的模块。模拟模块将为其访问的任何属性返回一个伪对象。因为我们是逐文件复制,依赖解析有时会发现模型文件导入但从未使用的文件(例如自定义序列化代码或训练助手)。使用此函数可以模拟这些功能,而无需修改原始代码。

参数:
  • include (Union[List[str], str]) –

    一个字符串,例如 "my_package.my_subpackage",或者要模拟的模块名称列表。字符串也可以是可能匹配多个模块的 glob 风格模式字符串。匹配此模式字符串的任何必需依赖项都将自动被模拟。

    示例

    'torch.**' – 匹配 torch 和 torch 的所有子模块,例如 'torch.nn''torch.nn.functional'

    'torch.*' – 匹配 'torch.nn''torch.functional',但不匹配 'torch.nn.functional'

  • exclude (Union[List[str], str]) – 一个可选模式,用于排除一些匹配 include 字符串的模式。例如 include='torch.**', exclude='torch.foo' 将模拟除 'torch.foo' 之外的所有 torch 包,默认值为 []

  • allow_empty (bool) – 一个可选标志,用于指定调用 mock() 方法时指定的模拟实现是否必须在打包期间与某个模块匹配。如果添加了 allow_empty=False 的 mock,并且调用了 close()(无论是显式调用还是通过 __exit__),而该 mock 未与被导出包所使用的模块匹配,则会抛出异常。如果 allow_empty=True,则不会抛出此类异常。

mocked_modules()[source]#

返回当前所有被模拟(mocked)的模块。

返回:

包含此包中将被模拟的模块名称的列表。

返回类型:

list[str]

register_extern_hook(hook)[source]#

在导出器上注册一个外部(extern)钩子。

每当一个模块与 extern() 模式匹配时,该钩子都会被调用。它应该具有以下签名

hook(exporter: PackageExporter, module_name: str) -> None

钩子将按注册顺序被调用。

返回:

一个句柄,可通过调用 handle.remove() 来移除已添加的钩子。

返回类型:

torch.utils.hooks.RemovableHandle

register_intern_hook(hook)[source]#

在导出器上注册一个内部(intern)钩子。

每当一个模块与 intern() 模式匹配时,该钩子都会被调用。它应该具有以下签名

hook(exporter: PackageExporter, module_name: str) -> None

钩子将按注册顺序被调用。

返回:

一个句柄,可通过调用 handle.remove() 来移除已添加的钩子。

返回类型:

torch.utils.hooks.RemovableHandle

register_mock_hook(hook)[source]#

在导出器上注册一个模拟(mock)钩子。

每当一个模块与 mock() 模式匹配时,该钩子都会被调用。它应该具有以下签名

hook(exporter: PackageExporter, module_name: str) -> None

钩子将按注册顺序被调用。

返回:

一个句柄,可通过调用 handle.remove() 来移除已添加的钩子。

返回类型:

torch.utils.hooks.RemovableHandle

save_binary(package, resource, binary)[source]#

将原始字节保存到包中。

参数:
  • package (str) – 该资源所属模块包的名称(例如 "my_package.my_subpackage")。

  • resource (str) – 资源的唯一名称,用于标识加载时的资源。

  • binary (str) – 要保存的数据。

save_module(module_name, dependencies=True)[source]#

module 的代码保存到包中。该模块的代码通过 importers 路径解析以找到模块对象,然后使用其 __file__ 属性找到源代码。

参数:
  • module_name (str) – 例如 my_package.my_subpackage,代码将被保存以提供该包的实现。

  • dependencies (bool, optional) – 如果为 True,我们将扫描源代码以获取依赖项。

save_pickle(package, resource, obj, dependencies=True, pickle_protocol=3)[source]#

使用 pickle 将 python 对象保存到归档文件中。等同于 torch.save(),但保存到归档文件中而不是独立文件。标准 pickle 不保存代码,仅保存对象。如果 dependencies 为真,此方法还将扫描被 pickle 的对象以确定重构它们所需的模块,并保存相关代码。

为了能够保存 type(obj).__name__my_module.MyObject 的对象,my_module.MyObject 必须能够根据 importer 顺序解析为该对象的类。当保存之前已打包的对象时,需要将导入器的 import_module 方法包含在 importer 列表中,此过程才能生效。

参数:
  • package (str) – 该资源所属模块包的名称(例如 "my_package.my_subpackage")。

  • resource (str) – 资源的唯一名称,用于标识加载时的资源。

  • obj (Any) – 要保存的对象,必须是可 pickle 的。

  • dependencies (bool, optional) – 如果为 True,我们将扫描源代码以获取依赖项。

save_source_file(module_name, file_or_directory, dependencies=True)[source]#

将本地文件系统中的 file_or_directory 添加到源包中,以提供 module_name 的代码。

参数:
  • module_name (str) – 例如 "my_package.my_subpackage",代码将被保存以提供该包的实现。

  • file_or_directory (str) – 代码文件或目录的路径。当传入目录时,目录中的所有 python 文件都将递归使用 save_source_file() 进行复制。如果文件名为 "/__init__.py",则该代码被视为一个包。

  • dependencies (bool, optional) – 如果为 True,我们将扫描源代码以获取依赖项。

save_source_string(module_name, src, is_package=False, dependencies=True)[source]#

src 作为导出包中 module_name 的源代码添加。

参数:
  • module_name (str) – 例如 my_package.my_subpackage,代码将被保存以提供该包的实现。

  • src (str) – 要为此包保存的 Python 源代码。

  • is_package (bool, optional) – 如果为 True,则将此模块视为一个包。包允许拥有子模块(例如 my_package.my_subpackage.my_subsubpackage),并且可以在其中保存资源。默认为 False

  • dependencies (bool, optional) – 如果为 True,我们将扫描源代码以获取依赖项。

save_text(package, resource, text)[source]#

将文本数据保存到包中。

参数:
  • package (str) – 该资源所属模块包的名称(例如 "my_package.my_subpackage")。

  • resource (str) – 资源的唯一名称,用于标识加载时的资源。

  • text (str) – 要保存的内容。

class torch.package.PackageImporter(file_or_buffer, module_allowed=<function PackageImporter.<lambda>>)[source]#

导入器允许您加载由 PackageExporter 写入包的代码。代码以密闭(hermetic)方式加载,使用包中的文件而不是标准的 python 导入系统。这使得可以对 PyTorch 模型代码和数据进行打包,以便在服务器上运行或用于以后的迁移学习。

包导入器确保模块中的代码只能从包内加载,除非是导出期间明确列为外部的模块。zip 归档文件中的 extern_modules 文件列出了包外部依赖的所有模块。这防止了“隐式”依赖,即包在本地运行时因为导入了本地安装的包而正常工作,但在复制到另一台机器时由于缺少依赖而失败。

__init__(file_or_buffer, module_allowed=<function PackageImporter.<lambda>>)[source]#

打开 file_or_buffer 以进行导入。这会检查导入的包是否仅需要 module_allowed 允许的模块。

参数:
  • file_or_buffer (str | PathLike[str] | IO[bytes] | PyTorchFileReader) – 类文件对象(必须实现 read(), readline(), tell()seek())、字符串,或包含文件名的 os.PathLike 对象。

  • module_allowed (Callable[[str], bool], optional) – 用于确定是否应允许外部提供的模块的方法。可用于确保加载的包不依赖于服务器不支持的模块。默认允许任何内容。

抛出:

ImportError – 如果包将使用不允许的模块。

file_structure(*, include='**', exclude=())[source]#

返回包 zip 文件的文件结构表示。

参数:
  • include (Union[List[str], str]) – 一个可选字符串,例如 "my_package.my_subpackage",或者包含在 zip 文件表示中的文件名列表。这也可以是 glob 风格的模式,如 PackageExporter.mock() 中所述。

  • exclude (Union[List[str], str]) – 一个可选模式,排除名称与该模式匹配的文件。

返回:

目录

返回类型:

目录

id()[source]#

返回 torch.package 用于区分 PackageImporter 实例的内部标识符。看起来像

<torch_package_0>
import_module(name, package=None)[source]#

如果模块尚未加载,则从包中加载模块,然后返回该模块。模块会加载到导入器本地,并将出现在 self.modules 中,而不是 sys.modules 中。

参数:
  • name (str) – 要加载的模块的完全限定名称。

  • package ([type], optional) – 未使用,但保留以匹配 importlib.import_module 的签名。默认为 None

返回:

(可能已加载的)模块。

返回类型:

types.ModuleType

load_binary(package, resource)[source]#

加载原始字节。

参数:
  • package (str) – 模块包的名称(例如 "my_package.my_subpackage")。

  • resource (str) – 资源的唯一名称。

返回:

加载的数据。

返回类型:

字节

load_pickle(package, resource, map_location=None)[source]#

从包中解封资源,使用 import_module() 加载构建对象所需的任何模块。

参数:
  • package (str) – 模块包的名称(例如 "my_package.my_subpackage")。

  • resource (str) – 资源的唯一名称。

  • map_location – 传递给 torch.load 以确定张量如何映射到设备。默认为 None

返回:

解封的对象。

返回类型:

任何

load_text(package, resource, encoding='utf-8', errors='strict')[source]#

加载字符串。

参数:
  • package (str) – 模块包的名称(例如 "my_package.my_subpackage")。

  • resource (str) – 资源的唯一名称。

  • encoding (str, optional) – 传递给 decode。默认为 'utf-8'

  • errors (str, optional) – 传递给 decode。默认为 'strict'

返回:

加载的文本。

返回类型:

str

python_version()[source]#

返回用于创建此包的 python 版本。

注意:此函数是实验性的且不具备前向兼容性。计划稍后将其移入锁文件(lock file)。

返回:

Optional[str] python 版本,例如 3.8.9;如果未随包存储版本,则为 None

class torch.package.Directory(name, is_dir)[source]#

一种文件结构表示。组织为 Directory 节点,其包含 Directory 子节点的列表。通过调用 PackageImporter.file_structure() 创建包的目录。

has_file(filename)[source]#

检查 Directory 中是否存在文件。

参数:

filename (str) – 要搜索的文件路径。

返回:

如果 Directory 包含指定文件。

返回类型:

布尔值