> ## Documentation Index
> Fetch the complete documentation index at: https://dragonwingdocs.qualcomm.com/llms.txt
> Use this file to discover all available pages before exploring further.

# 在 NPU 上运行 PyTorch 模型

您可以在 Dragonwing 设备的 NPU 上加速运行 PyTorch 模型,实现快速高效的机器学习推理。

## 量化并转换模型

PyTorch 模型的权重和激活值使用 32 位浮点数。而开发板上的 NPU 仅支持 8 位整数,因此必须对 PyTorch 模型进行\_量化\_ — 将浮点值转换为定点值。这会使模型更小、运行更快(并能够在 NPU 上运行),但会对精度产生一定影响。

在开发板上打开终端,或通过 SSH 连接到开发板,然后按照以下步骤操作:

1. 创建一个新的 `venv`,并安装 ONNX Runtime 以及示例所使用的 Python 软件包:

```shell theme={null}
mkdir -p ~/onnxruntime-demo/
cd ~/onnxruntime-demo/

python3.12 -m venv .venv
source .venv/bin/activate

# onnxruntime with AI Engine Direct bindings (only works on Python3.12)
wget https://cdn.edgeimpulse.com/qc-ai-docs/wheels/onnxruntime_qnn-1.23.0-cp312-cp312-linux_aarch64.whl
pip3 install onnxruntime_qnn-1.23.0-cp312-cp312-linux_aarch64.whl
rm onnxruntime*.whl

# Other dependencies
pip3 install Pillow onnx torch torchvision onnxscript
```

2. 以下是运行 torchvision 中 SqueezeNet-1.1 的端到端示例。将此文件保存为 `inference_pytorch_onnx.py`:

```python theme={null}
# pytorch_to_onnx_int8_squeezenet.py
import sys, os, glob, io, time, urllib.request, numpy as np, onnx, onnxruntime as ort
from PIL import Image

import torch, torchvision
import torchvision.transforms as T
from onnxruntime.quantization import (
    CalibrationDataReader,
    QuantFormat,
    QuantType,
    calibrate,
    quantize_static,
)

use_npu = True if len(sys.argv) >= 2 and sys.argv[1] == '--use-npu' else False

def download_file_if_not_exists(path, url):
    if not os.path.exists(path):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        print(f"Downloading {path} from {url}...")
        urllib.request.urlretrieve(url, path)
    return path

weights = torchvision.models.SqueezeNet1_1_Weights.DEFAULT
IMAGE_PATH = download_file_if_not_exists('images/boa-constrictor.jpg', 'https://cdn.edgeimpulse.com/qc-ai-docs/examples/boa-constrictor.jpg')

# Load PyTorch SqueezeNet1_1 from torchvision
device = "cpu"
model = torchvision.models.squeezenet1_1(weights)
model.eval().to(device)

# Export to ONNX (fp32)
os.makedirs("models", exist_ok=True)
onnx_fp32 = "models/squeezenet1_1_fp32.onnx"
input_size = 224
dummy = torch.randn(1, 3, input_size, input_size, device=device)

torch.onnx.export(
    model, dummy, onnx_fp32,
    input_names=["input"],
    output_names=["logits"],
    opset_version=13,
    do_constant_folding=True,
    # Keep the model shape static for QNNExecutionProvider / NPU execution.
)
onnx.checker.check_model(onnx.load(onnx_fp32))
print(f"Exported FP32 ONNX -> {onnx_fp32}")

# Provide a calibration data reader for static INT8 quantization
class ImageFolderDataReader(CalibrationDataReader):
    def __init__(self, image_paths):
        self.image_paths = image_paths
        self.transform = T.Compose([
            T.Resize(256), T.CenterCrop(224),
            T.ToTensor(), T.Normalize(mean=[0.485,0.456,0.406],
                                    std=[0.229,0.224,0.225])
        ])
        self._iter = None

    def get_next(self):
        if self._iter is None:
            self._iter = iter(self.image_paths)
        try:
            p = next(self._iter)
        except StopIteration:
            return None
        img = Image.open(p).convert("RGB")
        x = self.transform(img).unsqueeze(0).numpy()
        return {"input": x}

# Replace with representative images from your domain
calib = ImageFolderDataReader([IMAGE_PATH])

# Find the ONNX input name (matches "input" above, but we read it to be robust)
m = onnx.load(onnx_fp32)
onnx_input_name = m.graph.input[0].name

onnx_int8 = "models/squeezenet1_1_int8.onnx"

# Use QDQ format (widely supported); uint8 activations + int8 weights is a common choice
quantize_static(
    model_input=onnx_fp32,
    model_output=onnx_int8,
    calibration_data_reader=calib,
    activation_type=QuantType.QUInt8,
    weight_type=QuantType.QInt8
)
onnx.checker.check_model(onnx.load(onnx_int8))
print(f"Quantized INT8 ONNX -> {onnx_int8}")

# Use HTP backend of libQnnTFLiteDelegate.so (NPU) when --use-npu is passed in (otherwise CPU)
providers = []
if use_npu:
    providers.append(("QNNExecutionProvider", {
        "backend_type": "htp",
    }))
else:
    providers.append("CPUExecutionProvider")

input_shape = (1, 3, 224, 224)
img = Image.open(IMAGE_PATH).convert("RGB").resize((224, 224))
input_data = np.expand_dims(np.transpose(np.array(img, dtype=np.float32) / 255.0, (2, 0, 1)), 0)
#input_data = np.random.rand(*input_shape).astype(np.float32)

so = ort.SessionOptions()

sess = ort.InferenceSession(onnx_int8, sess_options=so, providers=providers)
actual_providers = sess.get_providers()
print(f"Using providers: {actual_providers}") # Show which providers are actually loaded

inputs  = sess.get_inputs()
outputs = sess.get_outputs()

_ = sess.run(None, { sess.get_inputs()[0].name: input_data })

# Run 10x so we can calculate avg. runtime per inference
start = time.perf_counter()
for i in range(10):
    out = sess.run(None, { sess.get_inputs()[0].name: input_data })
end = time.perf_counter()

def softmax(x, axis=-1):
    # subtract max for numerical stability
    x_max = np.max(x, axis=axis, keepdims=True)
    e_x = np.exp(x - x_max)
    return e_x / np.sum(e_x, axis=axis, keepdims=True)

scores = softmax(np.squeeze(out[0], axis=0))

# Take top 5
top_k_idx = scores.argsort()[-5:][::-1]

print("\nTop-5 predictions:")
for i in top_k_idx:
    label = weights.meta["categories"][i] if i < len(weights.meta["categories"]) else f"Class {i}"
    print(f"{label}: score={scores[i]}")

print("")
print(f'Inference took (on average): {((end - start) * 1000) / 10:.4g}ms. per image')
```

3. 在 CPU 上运行模型:

```shell theme={null}
python3 inference_pytorch_onnx.py

# Top-5 predictions:
# bassinet: score=0.473060667514801
# mosquito net: score=0.15856008231639862
# quilt: score=0.13602180778980255
# crib: score=0.07408633828163147
# cradle: score=0.03270549699664116
# 
# Inference took (on average): 7.418ms. per image
```

4. 在 NPU 上运行模型:

```shell theme={null}
python3 inference_pytorch_onnx.py --use-npu

# Top-5 predictions:
# bassinet: score=0.473060667514801
# mosquito net: score=0.15856008231639862
# quilt: score=0.13602180778980255
# crib: score=0.07408633828163147
# cradle: score=0.03270549699664116
# 
# Inference took (on average): 7.513ms. per image
```

## 技巧与窍门

### 禁用 CPU 回退

在调试时,您可能希望通过以下方式禁用回退到 CPU:

```shell theme={null}
so = ort.SessionOptions()
so.add_session_config_entry("session.disable_cpu_ep_fallback", "1")
```

### 构建新版本的 onnxruntime 软件包

参见:[edgeimpulse/onnxruntime-qnn-linux-aarch64](https://github.com/edgeimpulse/onnxruntime-qnn-linux-aarch64)。