下载
中文
注册

设置NPU为默认设备

TF Adapter提供了注册NPU设备的API,将NPU注册为TensorFlow的合法设备。开发者可以在带有@tf.function装饰并在CPU或GPU下可以正常工作的脚本文件开头,添加如下代码,设置NPU为默认设备。

1
2
3
import npu_device as npu
# 默认设置Device 0为计算设备
npu.open().as_default()

您应当在import其他python包前执行该操作,防止在加载后续包的过程中有未分发到NPU的算子执行行为。

npu.open接口的详细说明可参见npu.open