YOLOV8 兼容适配昇腾NPU
前提:
# 已经安装好对应显卡兼容的torch_npu插件,且通过拉取源码安装torchvision_npu
# 在项目启动入口上引入这俩个包
from transformers import is_torch_npu_available
if is_torch_npu_available():
import torch_npu
import torchvision_npu
基本代码兼容:
import torch
from transformers import is_torch_npu_available
# 在全局调用获取torch设备的时候,进行如下兼容,0可以切换,若有多个显卡的话
def get_yolo_device(cls):
"""
获取有效的设备进行训练
"""
# 多卡随机选择
device_count = device_id = 0
if torch.cuda