npu_print
函数功能
在图执行过程中,打印执行脚本中目标tensor值。该接口类似原生的print接口,但不会导致断图,详细功能介绍参见图内Tensor打印功能。
函数原型
def npu_print(*args, summarize_size=3)
参数说明
参数 |
输入/输出 |
说明 |
是否必选 |
||
---|---|---|---|---|---|
*args |
输入 |
位置入参,其中每个入参支持的数据类型为torch.Tensor、str、bool、float、int等基本类型,且其中至少包含一个Tensor类型输入。 |
是 |
||
summarize_size |
输入 |
tensor每个维度显示的数据元素个数,默认值为3。
例如shape为[10, 10]的tensor,当summarize_size=3,打印结果示例如下:
|
否 |
返回值说明
无
约束说明
- 打印为异步打印,打印输出顺序与图中执行顺序一致,与非图内执行的其他输出顺序无关。
例如图执行后,又在外部调用Python的print打印函数,可能出现图中打印位于Python print之后打印的情况。
- 接口为异步接口,打印数据会占用额外的Device内存同时耗费Device侧执行时间,设置过大的summarize_size可能会因内存不足或执行超时导致失败。
请合理设置summarize_size取值,通常适合打屏观察的数据量应该在KB级别以下。
例如shape为[10, 1024, 1024, 1024]的tensor,若设置summarize_size=3,内存占用量为KB级别;若设置summarize_size=1024,内存占用量会超过10G。
调用示例
1 2 3 4 5 6 7 8 9 |
import torch import torch_npu, torchair @torch.compile(backend="npu", fullgraph=True) def hello_tensor(x): torchair.ops.npu_print("hello, tensor:", x) v = torch.arange(10).npu() hello_tensor(v) # 打印结果为"hello, tensor: [0 1 2 ... 7 8 9]" |
父主题: ops接口