Torchvision Adapter插件 快速入门

运行环境变量

运行以下命令初始化CANN环境变量。

source /usr/local/Ascend/ascend-toolkit/set_env.sh

以上命令以root用户安装后的默认路径为例,请用户根据set_env.sh的实际路径进行替换。

Note

Docker中运行使用dvpp功能,需要映射/dev/dvpp_cmdlist设备文件,文件权限相关说明请参考对Docker进行安全加固

NPU适配

以Torchvision的torchvision.ops.nms算子为例,在CUDA/CPU环境中,该算子通过如下方法进行调用:

 # 算子的cuda/cpu版本调用
 import torch
 import torchvision
 
 ...
 torchvision.ops.nms(boxes, scores, iou_threshold) # boxes 和 scores 为 CPU/CUDA Tensor

安装Torchvision Adapter插件之后,只需增加import torchvision_npu,即可按照原生方式调用Torchvision算子。

 # 算子的npu版本调用
 import torch
 import torch_npu
 import torchvision
 import torchvision_npu
 
 ...
 torchvision.ops.nms(boxes, scores, iou_threshold) # boxes 和 scores 为 NPU Tensor