update: 适配昇腾

ultralytics/nn/autobackend.py

@@ -12,6 +12,10 @@ import numpy as np
 import torch
 import torch.nn as nn
 from PIL import Image
+import torch_npu
+from torch_npu.contrib import transfer_to_npu
+import torchair as tng
+from torchair.configs.compiler_config import CompilerConfig
 
 from ultralytics.utils import ARM64, IS_JETSON, IS_RASPBERRYPI, LINUX, LOGGER, PYTHON_VERSION, ROOT, yaml_load
 from ultralytics.utils.checks import check_requirements, check_suffix, check_version, check_yaml, is_rockchip
@@ -156,6 +160,12 @@ class AutoBackend(nn.Module):
             names = model.module.names if hasattr(model, "module") else model.names  # get class names
             model.half() if fp16 else model.float()
             self.model = model  # explicitly assign for to(), cpu(), cuda(), half()
+            # 添加torchair适配代码
+            config = CompilerConfig()
+            config.experimental_config.frozen_parameter = True
+            npu_backbend = tng.get_npu_backend(compiler_config=config)
+            model = torch.compile(model, dynamic=True, fullgraph=True, backend=npu_backbend)
+            tng.use_internal_format_weight(model.model)
             pt = True
 
         # PyTorch