Ubuntu 20.04 ARM64 YOLOv5完整实现指南

Ubuntu 20.04 ARM64 YOLOv5完整实现指南

一、环境配置

1.1 系统更新与基础依赖

# 更新系统包
sudo apt update && sudo apt upgrade -y

# 安装基础依赖
sudo apt install -y \
    git \
    wget \
    curl \
    vim \
    build-essential \
    cmake \
    pkg-config \
    libjpeg-dev \
    libtiff5-dev \
    libpng-dev \
    libavcodec-dev \
    libavformat-dev \
    libswscale-dev \
    libv4l-dev \
    libxvidcore-dev \
    libx264-dev \
    libgtk-3-dev \
    libatlas-base-dev \
    gfortran \
    python3-dev \
    python3-pip \
    python3-venv

1.2 Python环境配置

# 创建虚拟环境
python3 -m venv yolov5_env
source yolov5_env/bin/activate

# 升级pip
pip install --upgrade pip setuptools wheel

# 安装必要的Python包
pip install numpy scipy matplotlib ipython jupyter pandas sympy nose

1.3 PyTorch安装（ARM64专用）

# 对于ARM64架构，需要安装预编译的PyTorch版本
# 检查CUDA是否可用（如果有GPU）
python3 -c "import torch; print(torch.cuda.is_available())"

# 安装CPU版本的PyTorch（ARM64）
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

# 验证安装
python3 -c "import torch; print(torch.__version__); print(torch.cuda.is_available())"

1.4 OpenCV安装

# 方法1：pip安装（推荐）
pip install opencv-python opencv-python-headless

# 方法2：源码编译（如需要特定功能）
# git clone https://github.com/opencv/opencv.git
# cd opencv && mkdir build && cd build
# cmake -D CMAKE_BUILD_TYPE=RELEASE \
#       -D CMAKE_INSTALL_PREFIX=/usr/local \
#       -D OPENCV_EXTRA_MODULES_PATH=../../opencv_contrib/modules \
#       -D PYTHON_EXECUTABLE=$(which python3) ..
# make -j$(nproc)
# sudo make install

1.5 YOLOv5环境准备

# 克隆YOLOv5仓库
git clone https://github.com/ultralytics/yolov5.git
cd yolov5

# 安装YOLOv5依赖
pip install -r requirements.txt

# 安装额外依赖
pip install tensorboard seaborn thop

二、模型训练

2.1 数据准备

2.1.1 数据集结构

datasets/
└── custom_dataset/
    ├── images/
    │   ├── train/
    │   ├── val/
    │   └── test/
    └── labels/
        ├── train/
        ├── val/
        └── test/

2.1.2 创建数据集配置文件

# custom_dataset.yaml
path: /path/to/your/datasets/custom_dataset
train: images/train
val: images/val
test: images/test

nc: 3  # 类别数量
names: ['class1', 'class2', 'class3']  # 类别名称

2.1.3 标签格式转换脚本

# convert_labels.py
import os
import json
from pathlib import Path

def convert_coco_to_yolo(coco_json_path, output_dir):
    """将COCO格式转换为YOLO格式"""
    with open(coco_json_path, 'r') as f:
        data = json.load(f)
    
    # 创建类别ID映射
    category_mapping = {
   cat['id']: idx for idx, cat in enumerate(data['categories'])}
    
    # 获取图像信息
    images = {
   img['id']: img for img in data['images']}
    
    # 处理标注
    for ann in data['annotations']:
        img_id = ann['image_id']
        img_info = images[img_id]
        
        # 计算YOLO格式的坐标
        x, y, w, h = ann['bbox']
        img_w, img_h = img_info['width'], img_info['height']
        
        # 转换为相对坐标
        x_center = (x + w/2) / img_w
        y_center = (y + h/2) / img_h
        width = w / img_w
        height = h / img_h
        
        # 保存YOLO格式标签
        label_file = os.path.join(output_dir, f"{
     img_info['file_name'].split('.')[0]}.txt")
        with open(label_file, 'a') as f:
            f.write(f"{
     category_mapping[ann['category_id']]} {
     x_center} {
     y_center} {
     width} {
     height}\n")

# 使用示例
# convert_coco_to_yolo('annotations.json', 'labels/train/')

2.2 模型训练配置

2.2.1 创建训练脚本

# train_custom.py
import torch
import yaml
from yolov5.train import run

def main():
    # 训练参数配置
    opt = {
   
        'weights': 'yolov5s.pt',  # 预训练权重
        'cfg': '',  # 模型配置文件
        'data': 'custom_dataset.yaml',  # 数据集配置
        'hyp': 'data/hyps/hyp.scratch-low.yaml',  # 超参数配置
        'epochs': 100,  # 训练轮数
        'batch_size': 16,  # 批次大小（ARM64建议较小）
        'imgsz': 640,  # 输入图像尺寸
        'rect': False,  # 矩形训练
        'resume': False,  # 是否恢复训练
        'nosave': False,  # 是否保存模型
        'noval': False,  # 是否跳过验证
        'noautoanchor': False,  # 是否禁用自动锚点
        'evolve': None,  # 进化训练
        'bucket': '',  # gsutil bucket
        'cache': None,  # 图像缓存
        'image_weights'