torch.export¶
警告
此功能是正在积极开发中的原型,将来会有重大变化。
概述¶
torch.export.export()
接受一个任意的 Python 可调用对象(torch.nn.Module
、函数或方法),并生成一个跟踪图,该图以 Ahead-of-Time (AOT) 方式仅表示函数的张量计算,随后可以使用不同的输出执行或序列化。
import torch
from torch.export import export
class Mod(torch.nn.Module):
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
a = torch.sin(x)
b = torch.cos(y)
return a + b
example_args = (torch.randn(10, 10), torch.randn(10, 10))
exported_program: torch.export.ExportedProgram = export(
Mod(), args=example_args
)
print(exported_program)
ExportedProgram:
class GraphModule(torch.nn.Module):
def forward(self, arg0_1: f32[10, 10], arg1_1: f32[10, 10]):
# code: a = torch.sin(x)
sin: f32[10, 10] = torch.ops.aten.sin.default(arg0_1);
# code: b = torch.cos(y)
cos: f32[10, 10] = torch.ops.aten.cos.default(arg1_1);
# code: return a + b
add: f32[10, 10] = torch.ops.aten.add.Tensor(sin, cos);
return (add,)
Graph signature: ExportGraphSignature(
parameters=[],
buffers=[],
user_inputs=['arg0_1', 'arg1_1'],
user_outputs=['add'],
inputs_to_parameters={},
inputs_to_buffers={},
buffers_to_mutate={},
backward_signature=None,
assertion_dep_token=None,
)
Range constraints: {}
torch.export
生成一个干净的中间表示 (IR),具有以下不变性。有关 IR 的更多规范,请参见此处。
**健全性**:保证是原始程序的健全表示,并保持与原始程序相同的调用约定。
**规范化**:图中没有 Python 语义。原始程序中的子模块被内联以形成一个完全扁平化的计算图。
**图属性**:图是纯函数式的,这意味着它不包含具有副作用的操作,例如突变或别名。它不会改变任何中间值、参数或缓冲区。
**元数据**:图包含在跟踪过程中捕获的元数据,例如用户代码中的堆栈跟踪。
在底层,torch.export
利用了以下最新技术
**TorchDynamo (torch._dynamo)** 是一个内部 API,它使用名为帧评估 API 的 CPython 功能来安全地跟踪 PyTorch 图。这提供了极大改进的图捕获体验,需要重写的次数更少,以便完全跟踪 PyTorch 代码。
**AOT 自动微分** 提供了函数化的 PyTorch 图,并确保将图分解/降低到 ATen 运算符集。
**Torch FX (torch.fx)** 是图的底层表示形式,允许灵活的基于 Python 的转换。
现有框架¶
torch.compile()
也使用与 torch.export
相同的 PT2 堆栈,但略有不同
**JIT 与 AOT**:
torch.compile()
是一个 JIT 编译器,而它不打算用于在部署之外生成编译后的工件。**部分与完整图捕获**:当
torch.compile()
遇到模型中无法跟踪的部分时,它将“图中断”并回退到在 eager Python 运行时中运行程序。相比之下,torch.export
旨在获取 PyTorch 模型的完整图表示,因此当遇到无法跟踪的内容时,它会出错。由于torch.export
生成的完整图与任何 Python 功能或运行时分离,因此可以在不同的环境和语言中保存、加载和运行该图。**可用性权衡**:由于
torch.compile()
能够在遇到无法跟踪的内容时回退到 Python 运行时,因此它更加灵活。torch.export
则需要用户提供更多信息或重写其代码以使其可跟踪。
与 torch.fx.symbolic_trace()
相比,torch.export
使用 TorchDynamo 进行跟踪,TorchDynamo 在 Python 字节码级别运行,使其能够跟踪不受 Python 运算符重载支持的任意 Python 构造。此外,torch.export
精细地跟踪张量元数据,因此对张量形状等条件的条件语句不会导致跟踪失败。通常,torch.export
预计可用于更多用户程序,并生成更低级别的图(在 torch.ops.aten
运算符级别)。请注意,用户仍然可以在 torch.export
之前使用 torch.fx.symbolic_trace()
作为预处理步骤。
与 torch.jit.script()
相比,torch.export
不捕获 Python 控制流或数据结构,但它支持比 TorchScript 更多的 Python 语言功能(因为它更容易全面覆盖 Python 字节码)。生成的图更简单,并且只有直线控制流(显式控制流运算符除外)。
与 torch.jit.trace()
相比,torch.export
是健全的:它能够跟踪对大小执行整数计算的代码,并记录显示特定跟踪对其他输入有效的必要的所有边际条件。
导出 PyTorch 模型¶
示例¶
主要入口点是 torch.export.export()
,它接受一个可调用对象(torch.nn.Module
、函数或方法)和示例输入,并将计算图捕获到 torch.export.ExportedProgram
中。例如:
import torch
from torch.export import export
# Simple module for demonstration
class M(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
self.conv = torch.nn.Conv2d(
in_channels=3, out_channels=16, kernel_size=3, padding=1
)
self.relu = torch.nn.ReLU()
self.maxpool = torch.nn.MaxPool2d(kernel_size=3)
def forward(self, x: torch.Tensor, *, constant=None) -> torch.Tensor:
a = self.conv(x)
a.add_(constant)
return self.maxpool(self.relu(a))
example_args = (torch.randn(1, 3, 256, 256),)
example_kwargs = {"constant": torch.ones(1, 16, 256, 256)}
exported_program: torch.export.ExportedProgram = export(
M(), args=example_args, kwargs=example_kwargs
)
print(exported_program)
ExportedProgram:
class GraphModule(torch.nn.Module):
def forward(self, arg0_1: f32[16, 3, 3, 3], arg1_1: f32[16], arg2_1: f32[1, 3, 256, 256], arg3_1: f32[1, 16, 256, 256]):
# code: a = self.conv(x)
convolution: f32[1, 16, 256, 256] = torch.ops.aten.convolution.default(
arg2_1, arg0_1, arg1_1, [1, 1], [1, 1], [1, 1], False, [0, 0], 1
);
# code: a.add_(constant)
add: f32[1, 16, 256, 256] = torch.ops.aten.add.Tensor(convolution, arg3_1);
# code: return self.maxpool(self.relu(a))
relu: f32[1, 16, 256, 256] = torch.ops.aten.relu.default(add);
max_pool2d_with_indices = torch.ops.aten.max_pool2d_with_indices.default(
relu, [3, 3], [3, 3]
);
getitem: f32[1, 16, 85, 85] = max_pool2d_with_indices[0];
return (getitem,)
Graph signature: ExportGraphSignature(
parameters=['L__self___conv.weight', 'L__self___conv.bias'],
buffers=[],
user_inputs=['arg2_1', 'arg3_1'],
user_outputs=['getitem'],
inputs_to_parameters={
'arg0_1': 'L__self___conv.weight',
'arg1_1': 'L__self___conv.bias',
},
inputs_to_buffers={},
buffers_to_mutate={},
backward_signature=None,
assertion_dep_token=None,
)
Range constraints: {}
检查 ExportedProgram
,我们可以注意到以下几点:
torch.fx.Graph
包含原始程序的计算图,以及原始代码的记录,便于调试。该图仅包含在 此处 找到的
torch.ops.aten
算子和自定义算子,并且功能齐全,没有任何原地算子,例如torch.add_
。参数(卷积的权重和偏差)被提升为图的输入,导致图中没有
get_attr
节点,而这些节点之前存在于torch.fx.symbolic_trace()
的结果中。torch.export.ExportGraphSignature
对输入和输出签名进行建模,并指定哪些输入是参数。记录了图中每个节点生成的张量的结果形状和 dtype。例如,
convolution
节点将生成一个 dtype 为torch.float32
且形状为 (1, 16, 256, 256) 的张量。
非严格导出¶
在 PyTorch 2.3 中,我们引入了一种新的跟踪模式,称为**非严格模式**。它仍在不断完善中,因此如果您遇到任何问题,请使用“oncall: export”标签将其提交到 Github。
在*非严格模式*下,我们使用 Python 解释器跟踪程序。您的代码将像在 eager 模式下一样执行;唯一的区别是所有 Tensor 对象都将被 ProxyTensors 替换,ProxyTensors 会将它们的所有操作记录到一个图中。
在*严格模式*下(目前是默认模式),我们首先使用字节码分析引擎 TorchDynamo 跟踪程序。TorchDynamo 实际上并没有执行您的 Python 代码。相反,它会对其进行符号分析,并根据结果构建一个图。这种分析允许 torch.export 对安全性提供更强的保证,但并非所有 Python 代码都受支持。
例如,如果遇到一个不受支持的 TorchDynamo 功能,而该功能可能不容易解决,并且您知道 Python 代码对于计算来说并非绝对必要,那么您可能希望使用非严格模式。例如:
import contextlib
import torch
class ContextManager():
def __init__(self):
self.count = 0
def __enter__(self):
self.count += 1
def __exit__(self, exc_type, exc_value, traceback):
self.count -= 1
class M(torch.nn.Module):
def forward(self, x):
with ContextManager():
return x.sin() + x.cos()
export(M(), (torch.ones(3, 3),), strict=False) # Non-strict traces successfully
export(M(), (torch.ones(3, 3),)) # Strict mode fails with torch._dynamo.exc.Unsupported: ContextManager
在本例中,第一次使用非严格模式(通过 strict=False
标志)的调用成功地进行了跟踪,而第二次使用严格模式(默认)的调用则导致失败,因为 TorchDynamo 不支持上下文管理器。一种选择是重写代码(请参阅 torch.export 的限制),但由于上下文管理器不影响模型中的张量计算,我们可以使用非严格模式的结果。
表达动态性¶
默认情况下,torch.export
将跟踪程序,假设所有输入形状都是**静态**的,并将导出的程序专门用于这些维度。但是,某些维度(例如批处理维度)可以是动态的,并且在每次运行时都会发生变化。必须使用 torch.export.Dim()
API 创建此类维度,并通过 dynamic_shapes
参数将它们传递给 torch.export.export()
。例如:
import torch
from torch.export import Dim, export
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.branch1 = torch.nn.Sequential(
torch.nn.Linear(64, 32), torch.nn.ReLU()
)
self.branch2 = torch.nn.Sequential(
torch.nn.Linear(128, 64), torch.nn.ReLU()
)
self.buffer = torch.ones(32)
def forward(self, x1, x2):
out1 = self.branch1(x1)
out2 = self.branch2(x2)
return (out1 + self.buffer, out2)
example_args = (torch.randn(32, 64), torch.randn(32, 128))
# Create a dynamic batch size
batch = Dim("batch")
# Specify that the first dimension of each input is that batch size
dynamic_shapes = {"x1": {0: batch}, "x2": {0: batch}}
exported_program: torch.export.ExportedProgram = export(
M(), args=example_args, dynamic_shapes=dynamic_shapes
)
print(exported_program)
ExportedProgram:
class GraphModule(torch.nn.Module):
def forward(self, arg0_1: f32[32, 64], arg1_1: f32[32], arg2_1: f32[64, 128], arg3_1: f32[64], arg4_1: f32[32], arg5_1: f32[s0, 64], arg6_1: f32[s0, 128]):
# code: out1 = self.branch1(x1)
permute: f32[64, 32] = torch.ops.aten.permute.default(arg0_1, [1, 0]);
addmm: f32[s0, 32] = torch.ops.aten.addmm.default(arg1_1, arg5_1, permute);
relu: f32[s0, 32] = torch.ops.aten.relu.default(addmm);
# code: out2 = self.branch2(x2)
permute_1: f32[128, 64] = torch.ops.aten.permute.default(arg2_1, [1, 0]);
addmm_1: f32[s0, 64] = torch.ops.aten.addmm.default(arg3_1, arg6_1, permute_1);
relu_1: f32[s0, 64] = torch.ops.aten.relu.default(addmm_1); addmm_1 = None
# code: return (out1 + self.buffer, out2)
add: f32[s0, 32] = torch.ops.aten.add.Tensor(relu, arg4_1);
return (add, relu_1)
Graph signature: ExportGraphSignature(
parameters=[
'branch1.0.weight',
'branch1.0.bias',
'branch2.0.weight',
'branch2.0.bias',
],
buffers=['L__self___buffer'],
user_inputs=['arg5_1', 'arg6_1'],
user_outputs=['add', 'relu_1'],
inputs_to_parameters={
'arg0_1': 'branch1.0.weight',
'arg1_1': 'branch1.0.bias',
'arg2_1': 'branch2.0.weight',
'arg3_1': 'branch2.0.bias',
},
inputs_to_buffers={'arg4_1': 'L__self___buffer'},
buffers_to_mutate={},
backward_signature=None,
assertion_dep_token=None,
)
Range constraints: {s0: RangeConstraint(min_val=2, max_val=9223372036854775806)}
还有一些需要注意的事项:
通过
torch.export.Dim()
API 和dynamic_shapes
参数,我们将每个输入的第一个维度指定为动态的。查看输入arg5_1
和arg6_1
,它们的符号形状为 (s0, 64) 和 (s0, 128),而不是我们作为示例输入传入的 (32, 64) 和 (32, 128) 形状的张量。s0
是一个符号,表示此维度可以是一系列值。exported_program.range_constraints
描述了图中出现的每个符号的范围。在本例中,我们看到s0
的范围是 [2, inf]。由于技术原因(此处难以解释),假设它们不为 0 或 1。这不是错误,也不一定意味着导出的程序在维度为 0 或 1 时不起作用。有关此主题的深入讨论,请参阅 0/1 特化问题。
我们还可以指定输入形状之间更具表达力的关系,例如一对形状可能相差 1、一个形状可能是另一个形状的两倍或一个形状是偶数。例如:
class M(torch.nn.Module):
def forward(self, x, y):
return x + y[1:]
x, y = torch.randn(5), torch.randn(6)
dimx = torch.export.Dim("dimx", min=3, max=6)
dimy = dimx + 1
exported_program = torch.export.export(
M(), (x, y), dynamic_shapes=({0: dimx}, {0: dimy}),
)
print(exported_program)
ExportedProgram:
class GraphModule(torch.nn.Module):
def forward(self, arg0_1: "f32[s0]", arg1_1: "f32[s0 + 1]"):
# code: return x + y[1:]
slice_1: "f32[s0]" = torch.ops.aten.slice.Tensor(arg1_1, 0, 1, 9223372036854775807); arg1_1 = None
add: "f32[s0]" = torch.ops.aten.add.Tensor(arg0_1, slice_1); arg0_1 = slice_1 = None
return (add,)
Graph signature: ExportGraphSignature(
input_specs=[
InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='arg0_1'), target=None, persistent=None),
InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='arg1_1'), target=None, persistent=None)
],
output_specs=[
OutputSpec(kind=<OutputKind.USER_OUTPUT: 1>, arg=TensorArgument(name='add'), target=None)]
)
Range constraints: {s0: ValueRanges(lower=3, upper=6, is_bool=False), s0 + 1: ValueRanges(lower=4, upper=7, is_bool=False)}
需要注意的一些事项:
通过为第一个输入指定
{0: dimx}
,我们看到第一个输入的结果形状现在是动态的,即[s0]
。现在,通过为第二个输入指定{0: dimy}
,我们看到第二个输入的结果形状也是动态的。但是,由于我们表达了dimy = dimx + 1
,而不是arg1_1
的形状包含一个新符号,我们看到它现在用的是arg0_1
中使用的相同符号s0
。我们可以看到,dimy = dimx + 1
的关系通过s0 + 1
显示出来。查看范围约束,我们看到
s0
的范围是 [3, 6](最初指定),并且我们可以看到s0 + 1
的求解范围是 [4, 7]。
序列化¶
要保存 ExportedProgram
,用户可以使用 torch.export.save()
和 torch.export.load()
API。一种惯例是使用 .pt2
文件扩展名保存 ExportedProgram
。
例如:
import torch
import io
class MyModule(torch.nn.Module):
def forward(self, x):
return x + 10
exported_program = torch.export.export(MyModule(), torch.randn(5))
torch.export.save(exported_program, 'exported_program.pt2')
saved_exported_program = torch.export.load('exported_program.pt2')
特化¶
理解 torch.export
行为的一个关键概念是*静态*值和*动态*值之间的区别。
*动态*值是在每次运行时都可能发生变化的值。它们的行为类似于 Python 函数的普通参数,您可以为一个参数传递不同的值,并期望您的函数能够正确处理。张量*数据*被视为动态的。
*静态*值是在导出时固定的值,在导出的程序执行期间不能更改。当在跟踪过程中遇到该值时,导出器会将其视为常量并将其硬编码到图中。
当执行一个操作时(例如 x + y
)并且所有输入都是静态的,那么该操作的输出将直接硬编码到图中,并且该操作不会显示出来(即它将被常量折叠)。
当一个值被硬编码到图中时,我们说该图已经被*特化*为该值。
以下值是静态的:
输入张量形状¶
默认情况下,torch.export
将跟踪程序,并根据输入张量的形状进行特化,除非通过 torch.export
的 dynamic_shapes
参数将某个维度指定为动态的。这意味着,如果存在与形状相关的控制流,torch.export
将根据给定示例输入所采用的分支进行特化。例如:
import torch
from torch.export import export
class Mod(torch.nn.Module):
def forward(self, x):
if x.shape[0] > 5:
return x + 1
else:
return x - 1
example_inputs = (torch.rand(10, 2),)
exported_program = export(Mod(), example_inputs)
print(exported_program)
ExportedProgram:
class GraphModule(torch.nn.Module):
def forward(self, arg0_1: f32[10, 2]):
add: f32[10, 2] = torch.ops.aten.add.Tensor(arg0_1, 1);
return (add,)
(x.shape[0] > 5
) 的条件没有出现在 ExportedProgram
中,因为示例输入具有 (10, 2) 的静态形状。由于 torch.export
根据输入的静态形状进行特化,因此永远不会到达 else 分支 (x - 1
)。要保留跟踪图中基于张量形状的动态分支行为,需要使用 torch.export.dynamic_dim()
将输入张量 (x.shape[0]
) 的维度指定为动态的,并且需要 重写 源代码。
请注意,属于模块状态的张量(例如参数和缓冲区)始终具有静态形状。
Python 原语¶
torch.export
还专门针对 Python 原语,例如 int
、float
、bool
和 str
。但是,它们确实具有动态变体,例如 SymInt
、SymFloat
和 SymBool
。
例如:
import torch
from torch.export import export
class Mod(torch.nn.Module):
def forward(self, x: torch.Tensor, const: int, times: int):
for i in range(times):
x = x + const
return x
example_inputs = (torch.rand(2, 2), 1, 3)
exported_program = export(Mod(), example_inputs)
print(exported_program)
ExportedProgram:
class GraphModule(torch.nn.Module):
def forward(self, arg0_1: f32[2, 2], arg1_1, arg2_1):
add: f32[2, 2] = torch.ops.aten.add.Tensor(arg0_1, 1);
add_1: f32[2, 2] = torch.ops.aten.add.Tensor(add, 1);
add_2: f32[2, 2] = torch.ops.aten.add.Tensor(add_1, 1);
return (add_2,)
因为整数是专门化的,所以 torch.ops.aten.add.Tensor
操作都是使用硬编码常量 1
计算的,而不是 arg1_1
。如果用户在运行时为 arg1_1
传递了与导出时使用的值 1 不同的值(如 2),则会导致错误。此外,for
循环中使用的 times
迭代器也通过 3 次重复的 torch.ops.aten.add.Tensor
调用“内联”到图中,并且永远不会使用输入 arg2_1
。
Python 容器¶
Python 容器(List
、Dict
、NamedTuple
等)被认为具有静态结构。
torch.export 的限制¶
图中断¶
由于 torch.export
是一个从 PyTorch 程序捕获计算图的一次性过程,因此它最终可能会遇到程序中无法跟踪的部分,因为它几乎不可能支持跟踪所有 PyTorch 和 Python 功能。在 torch.compile
的情况下,不受支持的操作将导致“图中断”,并且不受支持的操作将使用默认的 Python 评估来运行。相比之下,torch.export
将要求用户提供其他信息或重写部分代码以使其可跟踪。由于跟踪基于在 Python 字节码级别进行评估的 TorchDynamo,因此与以前的跟踪框架相比,所需的重写要少得多。
当遇到图中断时,ExportDB 是一个很好的资源,用于了解支持和不支持的程序类型,以及重写程序以使其可跟踪的方法。
克服这些图中断的一种选择是使用 非严格导出。
数据/形状相关的控制流¶
当未专门化形状时,也可能在数据相关的控制流(if x.shape[0] > 2
)上遇到图中断,因为跟踪编译器不可能处理这种情况,而无需为组合爆炸数量的路径生成代码。在这种情况下,用户将需要使用特殊的控制流运算符来重写他们的代码。目前,我们支持 torch.cond 来表示类似 if-else 的控制流(即将推出更多!)。
运算符缺少 Fake/Meta/Abstract 内核¶
跟踪时,所有运算符都需要 FakeTensor 内核(也称为元内核、抽象实现)。这用于推理此运算符的输入/输出形状。
有关更多详细信息,请参阅 torch.library.register_fake()
。
如果您的模型使用了尚未实现 FakeTensor 内核的 ATen 运算符,请提交问题。
阅读更多¶
API 参考¶
- torch.export.export(mod, args, kwargs=None, *, dynamic_shapes=None, strict=True, preserve_module_call_signature=())[source]¶
export()
接收一个任意的 Python 可调用对象(一个 nn.Module、一个函数或一个方法)以及示例输入,并生成一个跟踪图,该图仅表示函数的张量计算,采用 Ahead-of-Time (AOT) 方式,随后可以使用不同的输入执行或序列化。跟踪图 (1) 在函数式 ATen 运算符集中生成规范化运算符(以及任何用户指定的自定义运算符),(2) 消除了所有 Python 控制流和数据结构(除某些例外),(3) 记录了形状约束集,以表明这种规范化和控制流消除对于未来的输入是合理的。健全性保证
在跟踪时,
export()
会记录用户程序和底层 PyTorch 运算符内核做出的与形状相关的假设。只有当这些假设成立时,输出ExportedProgram
才被认为是有效的。跟踪对输入张量的形状(而非值)进行假设。对于
export()
,此类假设必须在图捕获时进行验证才能成功。特别地:对输入张量的静态形状的假设会自动验证,无需额外的工作。
对输入张量的动态形状的假设需要通过使用
Dim()
API 构造动态维度并通过dynamic_shapes
参数将它们与示例输入相关联来明确指定。
如果任何假设无法验证,则会引发致命错误。发生这种情况时,错误消息将包括对规范的建议修复,这些修复是验证假设所必需的。例如,
export()
可能会建议对动态维度dim0_x
的定义进行以下修复,例如出现在与输入x
关联的形状中,该维度先前定义为Dim("dim0_x")
:dim = Dim("dim0_x", max=5)
此示例意味着生成的代码要求输入
x
的维度 0 小于或等于 5 才有效。您可以检查对动态维度定义的建议修复,然后将它们逐字复制到您的代码中,而无需更改对export()
调用的dynamic_shapes
参数。- 参数:
mod (Module) – 我们将跟踪此模块的 forward 方法。
dynamic_shapes (Optional[Union[Dict[str, Any], Tuple[Any], List[Any]]]) –
一个可选参数,其类型应为:1) 从
f
的参数名称到其动态形状规范的字典,2) 一个元组,按原始顺序指定每个输入的动态形状规范。如果您要指定关键字参数的动态性,则需要按照原始函数签名中定义的顺序传递它们。张量参数的动态形状可以指定为:(1) 从动态维度索引到
Dim()
类型的字典,其中不需要在此字典中包含静态维度索引,但当它们存在时,它们应映射到 None;或者 (2)Dim()
类型或 None 的元组/列表,其中Dim()
类型对应于动态维度,而静态维度由 None 表示。字典或张量元组/列表的参数通过使用包含规范的映射或序列递归指定。strict (bool) – 启用时(默认),导出函数将通过 TorchDynamo 跟踪程序,这将确保生成的图的正确性。否则,导出的程序将不会验证图中隐含的假设,并可能导致原始模型和导出模型之间的行为差异。当用户需要解决跟踪器中的错误,或者只是希望逐步启用模型中的安全性时,这很有用。请注意,这不会影响要生成的 IR 规范不同,并且模型将以相同的方式序列化,而不管此处传递了什么值。警告:此选项是实验性的,使用风险自负。
- 返回值
包含已跟踪可调用对象的
ExportedProgram
。- 返回类型
可接受的输入/输出类型
可接受的输入(对于
args
和kwargs
)和输出类型包括基本类型,即
torch.Tensor
、int
、float
、bool
和str
。数据类,但必须先通过调用
register_dataclass()
注册。包含
dict
、list
、tuple
、namedtuple
和OrderedDict
的(嵌套)数据结构,其中包含上述所有类型。
- torch.export.dynamic_shapes.dynamic_dim(t, index, debug_name=None)[source]¶
警告
(此功能已弃用。请参阅
Dim()
。)dynamic_dim()
构造一个_Constraint
对象,该对象描述张量t
的维度index
的动态性。_Constraint
对象应传递给export()
的constraints
参数。- 参数:
t (torch.Tensor) – 具有动态维度大小的示例输入张量
index (int) – 动态维度的索引
- 返回值
描述形状动态性的
_Constraint
对象。它可以传递给export()
,以便export()
不假设指定张量的大小是静态的,即将其保持为符号大小,而不是根据示例跟踪输入的大小进行专门化。
具体来说,
dynamic_dim()
可用于表示以下类型的动态性。维度的尺寸是动态且无界的
t0 = torch.rand(2, 3) t1 = torch.rand(3, 4) # First dimension of t0 can be dynamic size rather than always being static size 2 constraints = [dynamic_dim(t0, 0)] ep = export(fn, (t0, t1), constraints=constraints)
维度的尺寸是动态的,但有下界
t0 = torch.rand(10, 3) t1 = torch.rand(3, 4) # First dimension of t0 can be dynamic size with a lower bound of 5 (inclusive) # Second dimension of t1 can be dynamic size with a lower bound of 2 (exclusive) constraints = [ dynamic_dim(t0, 0) >= 5, dynamic_dim(t1, 1) > 2, ] ep = export(fn, (t0, t1), constraints=constraints)
维度的尺寸是动态的,但有上界
t0 = torch.rand(10, 3) t1 = torch.rand(3, 4) # First dimension of t0 can be dynamic size with a upper bound of 16 (inclusive) # Second dimension of t1 can be dynamic size with a upper bound of 8 (exclusive) constraints = [ dynamic_dim(t0, 0) <= 16, dynamic_dim(t1, 1) < 8, ] ep = export(fn, (t0, t1), constraints=constraints)
维度的尺寸是动态的,并且始终等于另一个动态维度的尺寸
t0 = torch.rand(10, 3) t1 = torch.rand(3, 4) # Sizes of second dimension of t0 and first dimension are always equal constraints = [ dynamic_dim(t0, 1) == dynamic_dim(t1, 0), ] ep = export(fn, (t0, t1), constraints=constraints)
混合和匹配上述所有类型,只要它们不表示冲突的要求
- torch.export.save(ep, f, *, extra_files=None, opset_version=None)[source]¶
警告
正在积极开发中,保存的文件可能无法在较新版本的 PyTorch 中使用。
将
ExportedProgram
保存到类文件对象。然后可以使用 Python APItorch.export.load
加载它。- 参数:
ep (ExportedProgram) – 要保存的已导出程序。
f (Union[str, os.PathLike, io.BytesIO) – 类文件对象(必须实现 write 和 flush)或包含文件名的字符串。
extra_files (Optional[Dict[str, Any]]) – 从文件名到内容的映射,这些内容将作为 f 的一部分存储。
示例
import torch import io class MyModule(torch.nn.Module): def forward(self, x): return x + 10 ep = torch.export.export(MyModule(), (torch.randn(5),)) # Save to file torch.export.save(ep, 'exported_program.pt2') # Save to io.BytesIO buffer buffer = io.BytesIO() torch.export.save(ep, buffer) # Save with extra files extra_files = {'foo.txt': b'bar'.decode('utf-8')} torch.export.save(ep, 'exported_program.pt2', extra_files=extra_files)
- torch.export.load(f, *, extra_files=None, expected_opset_version=None)[source]¶
警告
正在积极开发中,保存的文件可能无法在较新版本的 PyTorch 中使用。
加载之前使用
torch.export.save
保存的ExportedProgram
。- 参数:
ep (ExportedProgram) – 要保存的已导出程序。
f (Union[str, os.PathLike, io.BytesIO) – 类文件对象(必须实现 write 和 flush)或包含文件名的字符串。
extra_files (Optional[Dict[str, Any]]) – 此映射中提供的额外文件名将被加载,其内容将存储在提供的映射中。
expected_opset_version (Optional[Dict[str, int]]) – 操作集名称到预期操作集版本的映射
- 返回值
一个
ExportedProgram
对象- 返回类型
示例
import torch import io # Load ExportedProgram from file ep = torch.export.load('exported_program.pt2') # Load ExportedProgram from io.BytesIO object with open('exported_program.pt2', 'rb') as f: buffer = io.BytesIO(f.read()) buffer.seek(0) ep = torch.export.load(buffer) # Load with extra files. extra_files = {'foo.txt': ''} # values will be replaced with data ep = torch.export.load('exported_program.pt2', extra_files=extra_files) print(extra_files['foo.txt']) print(ep(torch.randn(5)))
- torch.export.register_dataclass(cls, *, serialized_type_name=None)[source]¶
将数据类注册为
torch.export.export()
的有效输入/输出类型。- 参数:
示例
@dataclass class InputDataClass: feature: torch.Tensor bias: int class OutputDataClass: res: torch.Tensor torch.export.register_dataclass(InputDataClass) torch.export.register_dataclass(OutputDataClass) def fn(o: InputDataClass) -> torch.Tensor: res = res=o.feature + o.bias return OutputDataClass(res=res) ep = torch.export.export(fn, (InputDataClass(torch.ones(2, 2), 1), )) print(ep)
- torch.export.dynamic_shapes.Dim(name, *, min=None, max=None)[source]¶
Dim()
构造一个类似于具有范围的命名符号整数的类型。它可以用于描述动态张量维度的多个可能值。请注意,相同张量或不同张量的不同动态维度可以用相同的类型来描述。
- class torch.export.dynamic_shapes.ShapesCollection[源代码]¶
dynamic_shapes 的构建器。用于为出现在输入中的张量分配动态形状规范。
- 示例:
args = ({“x”: tensor_x, “others”: [tensor_y, tensor_z]})
dim = torch.export.Dim(…) dynamic_shapes = torch.export.ShapesCollection() dynamic_shapes[tensor_x] = (dim, dim + 1, 8) dynamic_shapes[tensor_y] = {0: dim * 2} # 这等效于以下内容(现在自动生成): # dynamic_shapes = {“x”: (dim, dim + 1, 8), “others”: [{0: dim * 2}, None]}
torch.export(…, args, dynamic_shapes=dynamic_shapes)
- torch.export.dynamic_shapes.refine_dynamic_shapes_from_suggested_fixes(msg, dynamic_shapes)[源代码]¶
用于处理 export 的动态形状建议修复,和/或自动动态形状。根据 ConstraintViolation 错误消息和原始动态形状,优化给定的动态形状规范。
在大多数情况下,行为都很简单 - 例如,对于建议修复以专门化或优化 Dim 的范围,或建议派生关系的修复,新的动态形状规范将按此更新。
例如,建议修复
dim = Dim(‘dim’, min=3, max=6) -> 这只是优化了 dim 的范围 dim = 4 -> 这专门化为一个常量 dy = dx + 1 -> dy 被指定为一个独立的 dim,但实际上通过这种关系与 dx 绑定
但是,与派生 dims 相关的建议修复可能更加复杂。例如,如果为根 dim 提供了建议修复,则新派生的 dim 值将根据根进行评估。
例如,dx = Dim(‘dx’) dy = dx + 2 dynamic_shapes = {“x”: (dx,), “y”: (dy,)}
建议修复
dx = 4 # 特殊化将导致 dy 也特殊化为 6 dx = Dim(‘dx’, max=6) # dy 现在的 max = 8
派生的 dims 建议修复还可以用于表示可分性约束。这涉及创建未绑定到特定输入形状的新根 dims。在这种情况下,根 dims 不会直接出现在新的规范中,而是作为其中一个 dims 的根出现。
例如,建议修复
_dx = Dim(‘_dx’, max=1024) # 这不会出现在返回结果中,但 dx 会 dx = 4*_dx # dx 现在可以被 4 整除,最大值为 4096
- torch.export.Constraint¶
Union
[_Constraint
,_DerivedConstraint
] 的别名
- class torch.export.ExportedProgram(root, graph, graph_signature, state_dict, range_constraints, module_call_graph, example_inputs=None, verifier=None, tensor_constants=None, constants=None)[源代码]¶
来自
export()
的程序包。它包含一个表示张量计算的torch.fx.Graph
,一个包含所有提升参数和缓冲区张量值的 state_dict,以及各种元数据。您可以像调用由
export()
跟踪的原始可调用对象一样调用 ExportedProgram,使用相同的调用约定。要对图形执行转换,请使用
.module
属性访问torch.fx.GraphModule
。然后,您可以使用 FX 转换 重写图形。之后,您可以简单地再次使用export()
来构造正确的 ExportedProgram。- run_decompositions(decomp_table=None)[源代码]¶
对导出的程序运行一组分解,并返回一个新的导出程序。默认情况下,我们将运行核心 ATen 分解,以获取 核心 ATen 算子集 中的算子。
目前,我们不分解联合图。
- 返回类型
- class torch.export.ExportBackwardSignature(gradients_to_parameters: Dict[str, str], gradients_to_user_inputs: Dict[str, str], loss_output: str)[source]¶
- class torch.export.ExportGraphSignature(input_specs, output_specs)[source]¶
ExportGraphSignature
对导出图的输入/输出签名进行建模,导出图是一个具有更强不变性保证的 fx.Graph。导出图是函数式的,不会通过
getattr
节点访问图中的参数或缓冲区等“状态”。相反,export()
保证参数、缓冲区和常量张量作为输入从图中提取出来。同样,对缓冲区的任何更改也不包含在图中,而是将更改后的缓冲区的更新值建模为导出图的附加输出。所有输入和输出的顺序是
Inputs = [*parameters_buffers_constant_tensors, *flattened_user_inputs] Outputs = [*mutated_inputs, *flattened_user_outputs]
例如,如果导出以下模块
class CustomModule(nn.Module): def __init__(self): super(CustomModule, self).__init__() # Define a parameter self.my_parameter = nn.Parameter(torch.tensor(2.0)) # Define two buffers self.register_buffer('my_buffer1', torch.tensor(3.0)) self.register_buffer('my_buffer2', torch.tensor(4.0)) def forward(self, x1, x2): # Use the parameter, buffers, and both inputs in the forward method output = (x1 + self.my_parameter) * self.my_buffer1 + x2 * self.my_buffer2 # Mutate one of the buffers (e.g., increment it by 1) self.my_buffer2.add_(1.0) # In-place addition return output
生成的图将是
graph(): %arg0_1 := placeholder[target=arg0_1] %arg1_1 := placeholder[target=arg1_1] %arg2_1 := placeholder[target=arg2_1] %arg3_1 := placeholder[target=arg3_1] %arg4_1 := placeholder[target=arg4_1] %add_tensor := call_function[target=torch.ops.aten.add.Tensor](args = (%arg3_1, %arg0_1), kwargs = {}) %mul_tensor := call_function[target=torch.ops.aten.mul.Tensor](args = (%add_tensor, %arg1_1), kwargs = {}) %mul_tensor_1 := call_function[target=torch.ops.aten.mul.Tensor](args = (%arg4_1, %arg2_1), kwargs = {}) %add_tensor_1 := call_function[target=torch.ops.aten.add.Tensor](args = (%mul_tensor, %mul_tensor_1), kwargs = {}) %add_tensor_2 := call_function[target=torch.ops.aten.add.Tensor](args = (%arg2_1, 1.0), kwargs = {}) return (add_tensor_2, add_tensor_1)
生成的 ExportGraphSignature 将是
ExportGraphSignature( input_specs=[ InputSpec(kind=<InputKind.PARAMETER: 2>, arg=TensorArgument(name='arg0_1'), target='my_parameter'), InputSpec(kind=<InputKind.BUFFER: 3>, arg=TensorArgument(name='arg1_1'), target='my_buffer1'), InputSpec(kind=<InputKind.BUFFER: 3>, arg=TensorArgument(name='arg2_1'), target='my_buffer2'), InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='arg3_1'), target=None), InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='arg4_1'), target=None) ], output_specs=[ OutputSpec(kind=<OutputKind.BUFFER_MUTATION: 3>, arg=TensorArgument(name='add_2'), target='my_buffer2'), OutputSpec(kind=<OutputKind.USER_OUTPUT: 1>, arg=TensorArgument(name='add_1'), target=None) ] )
- class torch.export.ModuleCallSignature(inputs: List[Union[torch.export.graph_signature.TensorArgument, torch.export.graph_signature.SymIntArgument, torch.export.graph_signature.ConstantArgument, torch.export.graph_signature.CustomObjArgument, torch.export.graph_signature.TokenArgument]], outputs: List[Union[torch.export.graph_signature.TensorArgument, torch.export.graph_signature.SymIntArgument, torch.export.graph_signature.ConstantArgument, torch.export.graph_signature.CustomObjArgument, torch.export.graph_signature.TokenArgument]], in_spec: torch.utils._pytree.TreeSpec, out_spec: torch.utils._pytree.TreeSpec)[source]¶
- class torch.export.ModuleCallEntry(fqn: str, signature: Union[torch.export.exported_program.ModuleCallSignature, NoneType] = None)[source]¶
- class torch.export.graph_signature.InputSpec(kind: torch.export.graph_signature.InputKind, arg: Union[torch.export.graph_signature.TensorArgument, torch.export.graph_signature.SymIntArgument, torch.export.graph_signature.ConstantArgument, torch.export.graph_signature.CustomObjArgument, torch.export.graph_signature.TokenArgument], target: Union[str, NoneType], persistent: Union[bool, NoneType] = None)[source]¶
- class torch.export.graph_signature.OutputSpec(kind: torch.export.graph_signature.OutputKind, arg: Union[torch.export.graph_signature.TensorArgument, torch.export.graph_signature.SymIntArgument, torch.export.graph_signature.ConstantArgument, torch.export.graph_signature.CustomObjArgument, torch.export.graph_signature.TokenArgument], target: Union[str, NoneType])[source]¶
- class torch.export.graph_signature.ExportGraphSignature(input_specs, output_specs)[源代码]¶
ExportGraphSignature
为 Export Graph 的输入/输出签名建模,Export Graph 是一个具有更强不变性保证的 fx.Graph。Export Graph 是函数式的,并且不通过
getattr
节点访问图形内的“状态”,如参数或缓冲区。相反,export()
保证参数、缓冲区和常量张量作为输入从图形中提取出来。类似地,对缓冲区的任何修改也不包含在图形中,而是将已修改缓冲区的更新值建模为 Export Graph 的附加输出。所有输入和输出的顺序是
Inputs = [*parameters_buffers_constant_tensors, *flattened_user_inputs] Outputs = [*mutated_inputs, *flattened_user_outputs]
例如,如果导出以下模块
class CustomModule(nn.Module): def __init__(self): super(CustomModule, self).__init__() # Define a parameter self.my_parameter = nn.Parameter(torch.tensor(2.0)) # Define two buffers self.register_buffer('my_buffer1', torch.tensor(3.0)) self.register_buffer('my_buffer2', torch.tensor(4.0)) def forward(self, x1, x2): # Use the parameter, buffers, and both inputs in the forward method output = (x1 + self.my_parameter) * self.my_buffer1 + x2 * self.my_buffer2 # Mutate one of the buffers (e.g., increment it by 1) self.my_buffer2.add_(1.0) # In-place addition return output
生成的图将是
graph(): %arg0_1 := placeholder[target=arg0_1] %arg1_1 := placeholder[target=arg1_1] %arg2_1 := placeholder[target=arg2_1] %arg3_1 := placeholder[target=arg3_1] %arg4_1 := placeholder[target=arg4_1] %add_tensor := call_function[target=torch.ops.aten.add.Tensor](args = (%arg3_1, %arg0_1), kwargs = {}) %mul_tensor := call_function[target=torch.ops.aten.mul.Tensor](args = (%add_tensor, %arg1_1), kwargs = {}) %mul_tensor_1 := call_function[target=torch.ops.aten.mul.Tensor](args = (%arg4_1, %arg2_1), kwargs = {}) %add_tensor_1 := call_function[target=torch.ops.aten.add.Tensor](args = (%mul_tensor, %mul_tensor_1), kwargs = {}) %add_tensor_2 := call_function[target=torch.ops.aten.add.Tensor](args = (%arg2_1, 1.0), kwargs = {}) return (add_tensor_2, add_tensor_1)
生成的 ExportGraphSignature 将是
ExportGraphSignature( input_specs=[ InputSpec(kind=<InputKind.PARAMETER: 2>, arg=TensorArgument(name='arg0_1'), target='my_parameter'), InputSpec(kind=<InputKind.BUFFER: 3>, arg=TensorArgument(name='arg1_1'), target='my_buffer1'), InputSpec(kind=<InputKind.BUFFER: 3>, arg=TensorArgument(name='arg2_1'), target='my_buffer2'), InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='arg3_1'), target=None), InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='arg4_1'), target=None) ], output_specs=[ OutputSpec(kind=<OutputKind.BUFFER_MUTATION: 3>, arg=TensorArgument(name='add_2'), target='my_buffer2'), OutputSpec(kind=<OutputKind.USER_OUTPUT: 1>, arg=TensorArgument(name='add_1'), target=None) ] )
- class torch.export.unflatten.InterpreterModule(graph)[源代码]¶
一个使用 torch.fx.Interpreter 执行而不是 GraphModule 使用的通常代码生成的模块。这提供了更好的堆栈跟踪信息,并使调试执行更容易。
- torch.export.unflatten.unflatten(module, flat_args_adapter=None)[源代码]¶
展平 ExportedProgram,生成一个与原始 Eager 模块具有相同模块层次结构的模块。如果您尝试将
torch.export
与另一个需要模块层次结构而不是torch.export
通常生成的平面图的系统一起使用,这将非常有用。注意
展平模块的参数/关键字参数不一定与 Eager 模块匹配,因此进行模块交换(例如
self.submod = new_mod
)不一定有效。如果您需要交换模块,则需要设置torch.export.export()
的preserve_module_call_signature
参数。- 参数:
**module** (ExportedProgram) – 要展平的 ExportedProgram。
**flat_args_adapter** (可选[FlatArgsAdapter]) – 如果输入 TreeSpec 与导出的模块不匹配,则调整平面参数。
- 返回值
UnflattenedModule
的实例,它与导出前的原始 Eager 模块具有相同的模块层次结构。- 返回类型
UnflattenedModule