用于导出TorchAir生成的离线图(air格式),导出的图不再依赖PyTorch框架,可直接由CANN软件栈加载执行,减少框架调度带来的性能损耗,方便在不同的部署环境上移植,功能详情参见Dynamo导图功能。
def dynamo_export(*args, model: torch.nn.Module, export_path: str = "export_file", export_name: str = "export", dynamic: bool = False, config=CompilerConfig(), **kwargs)
参数 |
输入/输出 |
说明 |
是否必选 |
---|---|---|---|
model |
输入 |
需要导出的模型,继承原生的torch.nn.Module类。 |
是 |
export_path |
输入 |
离线图导出的文件存放路径,字符串类型,缺省值为当前路径下的"export_file"文件。 |
否 |
export_name |
输入 |
离线图导出的图名称,字符串类型,缺省值为"export"。 |
否 |
dynamic |
输入 |
设置导出静态模型还是动态模型,布尔类型。
|
否 |
config |
输入 |
图编译配置,CompilerConfig类的实例化,缺省情况下采用TorchAir自动生成的配置。 导图时支持配置auto_atc_config_generated和enable_record_nn_module_stack功能,配置样例参见使用方法。 |
否 |
*args、**kwargs |
输入 |
导出模型时的样例输入,不同的输入可能导致模型走入不同的分支,进而导致trace的图不同。应当选取执行推理时的典型值。 |
否 |
无
以单卡场景下的导图过程为例,代码如下,导图结果和多卡场景下的导图示例请参见使用示例。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 | import torch import torch_npu import torchair class Model(torch.nn.Module): def __init__(self): super().__init__() self.p1 = torch.nn.Parameter(torch.randn(2, 4)) self.p2 = torch.nn.Parameter(torch.randn(2, 4)) def forward(self, x, y): x = x + y + self.p1 + self.p2 return x model = Model() x = torch.randn(2, 4) y = torch.randn(2, 4) torchair.dynamo_export(x, y, model=model, export_path="./test_export_file_False", dynamic=False) |