RF-DETR模型迁移指导文档

cd /opt
git clone https://github.com/roboflow/rf-detr.git

python3.10 -m venv /opt/rfdetr_venv

export LD_LIBRARY_PATH=/usr/local/python3.10/lib:$LD_LIBRARY_PATH
export PYTHONPATH=/opt/rfdetr_venv/lib/python3.10/site-packages:/rf-detr/src:$PYTHONPATH
export PYTHONNOUSERSITE=1
export TASK_QUEUE_ENABLE=0
export ACL_OP_INIT_MODE=1
export TBE_PARALLEL_COMPILER=0
source /usr/local/Ascend/ascend-toolkit/8.2.RC1/aarch64-linux/script/set_env.sh
source /usr/local/Ascend/nnal/atb/set_env.sh

pip install torch==2.4.0
pip install torch_npu==2.4.0
pip install transformers peft pydantic supervision tqdm requests
pip install pycocotools scipy albumentations

pip install attrs decorator
cp /opt/rfdetr_venv/lib/python3.10/site-packages/attrs /usr/local/Ascend/ascend-toolkit/8.2.RC1/python/site-packages/ -r
cp /opt/rfdetr_venv/lib/python3.10/site-packages/decorator.py /usr/local/Ascend/ascend-toolkit/8.2.RC1/python/site-packages/

/data/coco_formatted/
├── annotations/
│   ├── instances_train2017.json
│   └── instances_val2017.json
├── train2017/
│   └── *.jpg
└── val2017/
    └── *.jpg

export ASCEND_RT_VISIBLE_DEVICES=4
python train_npu_native.py --epochs 10 --batch-size 8 --output-dir output/rfdetr_npu_single

./run_ddp.sh

export ASCEND_RT_VISIBLE_DEVICES=4,5,6,7
/usr/local/python3.10/bin/python3.10 -m torch.distributed.launch \
    --nproc_per_node=4 \
    --master_addr=127.0.0.1 \
    --master_port=29500 \
    --use_env \
    train_npu_ddp.py \
    --epochs 50 \
    --batch-size 8 \
    --output-dir output/rfdetr_npu_4card

Epoch 1 completed, loss: 11.1578
Epoch 10 completed, loss: 8.2097
Epoch 20 completed, loss: 7.7716
Epoch 30 completed, loss: 7.5393
Epoch 40 completed, loss: 7.3361
Epoch 50 completed, loss: 6.2836
Training completed! Best loss: 6.2836

pip install onnx

import os
import sys

os.environ['TASK_QUEUE_ENABLE'] = '0'
os.environ['ACL_OP_INIT_MODE'] = '1'
os.environ['TBE_PARALLEL_COMPILER'] = '0'

venv_site_packages = "/opt/rfdetr_venv/lib/python3.10/site-packages"
sys.path.insert(0, venv_site_packages)
sys.path = [p for p in sys.path if 'python3.11' not in p]

import torch

from rfdetr.config import RFDETRNanoConfig, TrainConfig
from rfdetr._namespace import build_namespace
from rfdetr.models.lwdetr import build_model

model_config = RFDETRNanoConfig(num_classes=80)
train_config = TrainConfig(dataset_dir="/data/coco_formatted")
ns = build_namespace(model_config, train_config)

model = build_model(args=ns)

checkpoint = torch.load('best_model.pth', map_location='cpu')
model.load_state_dict(checkpoint['model_state_dict'])
model = model.float()
model.eval()

print("Calling model.export() to prepare for ONNX export...")
model.export()

resolution = 384
dummy_input = torch.randn(1, 3, resolution, resolution).float()

print(f"Exporting ONNX with resolution {resolution}x{resolution}...")

torch.onnx.export(
    model,
    dummy_input,
    'best_model.onnx',
    export_params=True,
    opset_version=17,
    do_constant_folding=True,
    input_names=['images'],
    output_names=['output'],
    dynamic_axes=None,
    dynamo=False
)

print("转换完成！输出文件: best_model.onnx")

export LD_LIBRARY_PATH=/usr/local/python3.10/lib:$LD_LIBRARY_PATH
export PYTHONPATH=/opt/rfdetr_venv/lib/python3.10/site-packages:/rf-detr/src:$PYTHONPATH
/opt/rfdetr_venv/bin/python3.10 export_onnx.py

python -c "import onnx; model = onnx.load('best_model.onnx'); onnx.checker.check_model(model); print('ONNX模型验证通过！')"

#!/bin/bash

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

export TBE_PARALLEL_COMPILER=0
export ASCEND_SLOG_PRINT_TO_STDOUT=0
export ASCEND_GLOBAL_LOG_LEVEL=3

export CPLUS_INCLUDE_PATH=/usr/include/c++/12:/usr/include/c++/12/aarch64-openEuler-linux:$CPLUS_INCLUDE_PATH
export C_INCLUDE_PATH=/usr/include:$C_INCLUDE_PATH

ONNX_MODEL="best_model.onnx"
OM_MODEL="best_model"
CHIP_TYPE="Ascend910B3"
INPUT_SHAPE="images:1,3,384,384"
PRECISION="allow_mix_precision"

echo "开始转换ONNX模型到OM格式..."
atc --model=$ONNX_MODEL \
    --framework=5 \
    --output=$OM_MODEL \
    --input_format=NCHW \
    --input_shape=$INPUT_SHAPE \
    --precision_mode=$PRECISION \
    --soc_version=$CHIP_TYPE

if [ $? -eq 0 ]; then
    echo "模型转换成功！OM文件已保存为: $OM_MODEL.om"
else
    echo "模型转换失败，请检查错误信息。"
    exit 1
fi

chmod +x convert_onnx_to_om.sh
./convert_onnx_to_om.sh

pip install ais_bench

import cv2
import numpy as np
import glob
import os
import random

os.makedirs("bin", exist_ok=True)

MEAN = np.array([0.485, 0.456, 0.406])
STD = np.array([0.229, 0.224, 0.225])

val_images = glob.glob("/data/coco_formatted/val/*.jpg")
selected_images = random.sample(val_images, min(10, len(val_images)))

for f in selected_images:
    img = cv2.imread(f)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, (384, 384))
    img = img.astype(np.float32) / 255.0
    img = (img - MEAN) / STD
    img = img.transpose(2, 0, 1)[None, ...].astype(np.float32)
    bin_name = f"bin/" + os.path.basename(f)[:-4] + ".bin"
    img.tofile(bin_name)
    print(f"Converted: {f} -> {bin_name}")

print(f"\nTotal converted: {len(selected_images)} images")

python3 jpg_bin.py

import os
import sys
import glob
import numpy as np

sys.path.insert(0, '/usr/local/lib/python3.11/site-packages')
sys.path.insert(0, '/usr/local/lib64/python3.11/site-packages')

from ais_bench.infer.interface import InferSession

os.makedirs("results", exist_ok=True)

session = InferSession(device_id=0, model_path="best_model.om")

bin_files = glob.glob("bin/*.bin")
print(f"Found {len(bin_files)} bin files")

for bin_file in bin_files:
    input_data = np.fromfile(bin_file, dtype=np.float32).reshape(1, 3, 384, 384)
    
    outputs = session.infer([input_data])
    
    result_name = f"results/" + os.path.basename(bin_file)[:-4] + "_output.bin"
    
    if isinstance(outputs, list):
        for i, output in enumerate(outputs):
            output_name = f"results/" + os.path.basename(bin_file)[:-4] + f"_output_{i}.bin"
            if isinstance(output, np.ndarray):
                output.tofile(output_name)
                print(f"Saved: {output_name}, shape: {output.shape}")
    else:
        if isinstance(outputs, np.ndarray):
            outputs.tofile(result_name)
            print(f"Saved: {result_name}, shape: {outputs.shape}")
    
    print(f"Processed: {bin_file}")

print(f"\nTotal processed: {len(bin_files)} images")

export ASCEND_RT_VISIBLE_DEVICES=0
python3 om_inference.py

source /usr/local/Ascend/ascend-toolkit/8.2.RC1/aarch64-linux/script/set_env.sh
python3.11 inference_om.py \
    -m best_model_epoch248.om \
    -i /data/coco_formatted/val/000000200839.jpg \
    --threshold 0.5 \
    --output result.jpg

Detected 6 objects (threshold=0.5)
  [1] airplane: score=0.9427, box=[92.2, 170.1, 477.2, 316.0]
  [2] refrigerator: score=0.7923, box=[142.9, 240.4, 179.3, 271.9]
  [3] refrigerator: score=0.7400, box=[199.4, 237.6, 230.6, 275.2]
  [4] refrigerator: score=0.7370, box=[234.5, 201.8, 265.8, 244.2]
  [5] refrigerator: score=0.6601, box=[360.1, 193.9, 408.7, 260.4]
  [6] refrigerator: score=0.5389, box=[305.8, 192.2, 341.7, 243.2]