PyTorch-2.x镜像结合TPH-YOLOv5的完整部署流程

1. 引言：为什么选择PyTorch-2.x通用开发镜像？

在深度学习项目中，环境配置往往是第一步也是最繁琐的一步。尤其是当我们要部署一个复杂的模型如TPH-YOLOv5时，依赖冲突、CUDA版本不匹配、包缺失等问题常常让人头疼。

本文将带你使用一款开箱即用的深度学习镜像——PyTorch-2.x-Universal-Dev-v1.0，完成从环境启动到TPH-YOLOv5模型训练与推理的全流程部署。这款镜像基于官方PyTorch构建，预装了常用数据处理和可视化工具，并已配置国内源，极大提升了安装效率。

我们将聚焦于无人机场景下的目标检测任务，利用TPH-YOLOv5在VisDrone数据集上的优异表现，展示如何在一个纯净、高效、稳定的环境中快速落地先进模型。

1.1 TPH-YOLOv5是什么？它解决了什么问题？

TPH-YOLOv5是YOLOv5的一个改进版本，专为无人机航拍图像中的目标检测设计。相比传统YOLOv5，它针对以下三大挑战进行了优化：

目标尺度变化剧烈：无人机飞行高度不同导致物体大小差异巨大。
高密度目标遮挡严重：城市或交通场景中车辆、行人密集。
背景复杂且干扰多：大范围地理覆盖带来大量无关信息。

为此，TPH-YOLOv5引入了三项关键技术：

增加一个预测头用于检测微小物体；
使用Transformer Prediction Heads (TPH)替代原生卷积头，提升对上下文的理解能力；
集成CBAM注意力模块，帮助网络聚焦关键区域。

最终，在VisDrone2021测试集上达到39.18% mAP，超越前SOTA方法1.81%，位列挑战赛第五名。

2. 环境准备：启动并验证PyTorch-2.x镜像

2.1 启动镜像环境

假设你正在使用支持容器化AI开发平台（如CSDN星图、Docker或Kubernetes），可通过以下命令拉取并运行该镜像：

docker run -it --gpus all \ -p 8888:8888 \ -v ./tph-yolov5-workspace:/workspace \ pytorch-universal-dev:v1.0

注：实际镜像名称请根据平台替换为PyTorch-2.x-Universal-Dev-v1.0。

该镜像特点如下：

Python 3.10+
支持 CUDA 11.8 / 12.1，兼容RTX 30/40系列及A800/H800
已预装 JupyterLab、NumPy、Pandas、OpenCV、Matplotlib 等常用库
默认配置阿里云/清华源，pip install速度显著提升

2.2 验证GPU与PyTorch可用性

进入容器后，首先确认GPU是否正常挂载：

nvidia-smi

输出应显示你的显卡型号和驱动状态。

接着检查PyTorch是否能识别CUDA：

import torch print("CUDA available:", torch.cuda.is_available()) print("CUDA version:", torch.version.cuda) print("Number of GPUs:", torch.cuda.device_count()) print("Current GPU:", torch.cuda.get_device_name(0))

预期输出示例：

CUDA available: True CUDA version: 11.8 Number of GPUs: 1 Current GPU: NVIDIA GeForce RTX 3090

如果以上均通过，则说明环境已就绪，可以开始部署TPH-YOLOv5。

3. 模型部署：获取并配置TPH-YOLOv5代码库

3.1 克隆项目仓库

TPH-YOLOv5通常基于YOLOv5代码库进行修改。我们从原始YOLOv5仓库克隆基础代码，并切换至合适分支：

git clone https://github.com/ultralytics/yolov5.git cd yolov5 git checkout v7.0 # 推荐稳定版本

然后下载TPH-YOLOv5的补丁文件或专用分支（假设由作者提供）：

wget https://example.com/tph-yolov5-patch.zip unzip tph-yolov5-patch.zip cp -r patch/* .

实际操作中，请参考论文作者公开的GitHub仓库地址。

3.2 安装必要依赖

虽然镜像已包含大部分常用包，但仍需安装YOLOv5特定依赖：

pip install -r requirements.txt

由于镜像已配置国内源，此过程通常只需1-2分钟即可完成。

3.3 目录结构说明

标准YOLOv5项目结构如下：

yolov5/ ├── models/ # 模型定义（含.yaml配置） ├── utils/ # 工具函数 ├── data/ # 数据集配置文件 ├── weights/ # 存放预训练权重 ├── runs/ # 训练日志与结果保存路径 └── train.py, detect.py # 主程序入口

我们需要在models/中添加TPH-YOLOv5的自定义模块。

4. 模型实现：集成TPH与CBAM模块

4.1 添加Transformer Prediction Head (TPH)

TPH的核心是在检测头中引入Transformer编码器块，以增强对全局上下文的感知能力。

编辑models/common.py，新增TPH模块：

import torch import torch.nn as nn from torch.nn import MultiheadAttention class TransformerPredictHead(nn.Module): def __init__(self, c_in, num_heads=8): super().__init__() self.attention = MultiheadAttention(embed_dim=c_in, num_heads=num_heads) self.norm = nn.LayerNorm(c_in) self.ffn = nn.Sequential( nn.Linear(c_in, c_in * 4), nn.ReLU(), nn.Linear(c_in * 4, c_in) ) def forward(self, x): # x: [B, C, H, W] -> [H*W, B, C] b, c, h, w = x.shape x = x.view(b, c, -1).permute(2, 0, 1) # [N, B, C] x_att, _ = self.attention(x, x, x) x = self.norm(x + x_att) x = self.norm(x + self.ffn(x)) x = x.permute(1, 2, 0).view(b, c, h, w) return x

4.2 在YOLOv5头部替换为TPH

修改models/yolo.py中的Detect类，在初始化时插入TPH模块：

class Detect(nn.Module): def __init__(self, nc=80, anchors=(), ch=(), inplace=True): super().__init__() self.nc = nc self.no = nc + 5 self.nl = len(anchors) self.na = len(anchors[0]) // 2 self.grid = [torch.zeros(1)] * self.nl self.anchor_grid = [torch.zeros(1)] * self.nl self.stride = None # 增加TPH模块 self.tph = nn.ModuleList(TransformerPredictHead(x) for x in ch)

并在前向传播中调用：

def forward(self, x): z = [] for i in range(self.nl): x[i] = self.tph[i](x[i]) # 应用TPH ... return x, z

4.3 集成CBAM注意力机制

CBAM模块可沿通道和空间两个维度生成注意力图，提升对复杂背景的鲁棒性。

创建utils/cbam.py：

class ChannelAttention(nn.Module): def __init__(self, in_planes, ratio=16): super().__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.max_pool = nn.AdaptiveMaxPool2d(1) self.fc = nn.Sequential(nn.Conv2d(in_planes, in_planes // ratio, 1), nn.ReLU(), nn.Conv2d(in_planes // ratio, in_planes, 1)) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = self.fc(self.avg_pool(x)) max_out = self.fc(self.max_pool(x)) return self.sigmoid(avg_out + max_out) class SpatialAttention(nn.Module): def __init__(self, kernel_size=7): super().__init__() self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size//2) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = torch.mean(x, dim=1, keepdim=True) max_out, _ = torch.max(x, dim=1, keepdim=True) x = torch.cat([avg_out, max_out], dim=1) return self.sigmoid(self.conv(x)) class CBAM(nn.Module): def __init__(self, c_in, ratio=16, kernel_size=7): super().__init__() self.ca = ChannelAttention(c_in, ratio) self.sa = SpatialAttention(kernel_size) def forward(self, x): x = x * self.ca(x) x = x * self.sa(x) return x

将其插入骨干网络末端（如CSPDarknet53之后）：

# 在 models/common.py 的 Focus 或 Conv 后添加 self.cbam = CBAM(channels)

5. 数据准备：VisDrone2021数据集处理

5.1 下载与解压数据集

VisDrone2021是一个大规模无人机航拍目标检测数据集，包含训练、验证和测试集。

mkdir -p datasets/visdrone cd datasets/visdrone wget http://aiskyeye.com/download/VisDrone2021-DET-train.zip wget http://aiskyeye.com/download/VisDrone2021-DET-val.zip unzip VisDrone2021-DET-train.zip unzip VisDrone2021-DET-val.zip

5.2 转换标注格式

VisDrone使用.txt格式标注，每行表示一个对象：

<x_left> <y_top> <width> <height> <score> <object_category> <occlusion> <truncation>

我们需要提取前六列，并转换为YOLO格式（归一化中心坐标）：

import os def convert_visdrone_to_yolo(txt_path, img_w, img_h): with open(txt_path, 'r') as f: lines = f.readlines() yolo_lines = [] for line in lines: parts = line.strip().split(',') if len(parts) < 6: continue category = int(parts[5]) if category == 0 or category > 10: # 忽略无效类别 continue x_center = (float(parts[0]) + float(parts[2]) / 2) / img_w y_center = (float(parts[1]) + float(parts[3]) / 2) / img_h w = float(parts[2]) / img_w h = float(parts[3]) / img_h yolo_lines.append(f"{category-1} {x_center} {y_center} {w} {h}") return yolo_lines

批量处理所有标签文件并保存至labels/目录。

5.3 创建数据配置文件

新建data/visdrone.yaml：

train: ../datasets/visdrone/VisDrone2021-DET-train/images val: ../datasets/visdrone/VisDrone2021-DET-val/images nc: 10 names: ['pedestrian', 'people', 'bicycle', 'car', 'van', 'truck', 'tricycle', 'awning-tricycle', 'bus', 'motor']

6. 模型训练：启动TPH-YOLOv5训练任务

6.1 准备预训练权重

使用YOLOv5x作为基线模型，加载其预训练权重有助于加速收敛：

wget https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5x.pt -P weights/

6.2 修改模型配置文件

编辑models/yolov5x-tph.yaml：

# Parameters nc: 10 depth_multiple: 1.33 width_multiple: 1.25 # YOLOv5 backbone backbone: - [-1, 1, Conv, [128, 3, 2]] # 0-P1/2 ... - [-1, 3, C3, [1024, True]] # 23 # 添加CBAM - [-1, 1, CBAM, [1024]] # TPH部分将在Detect层自动应用 # YOLOv5 head head: - [-1, 1, nn.Upsample, [None, 2, 'nearest']] ... - [-1, 1, Detect, [nc, anchors]]

6.3 启动训练

python train.py \ --img 1536 \ --batch 2 \ --epochs 65 \ --data data/visdrone.yaml \ --weights weights/yolov5x.pt \ --cfg models/yolov5x-tph.yaml \ --name tph-yolov5-visdrone \ --device 0

注意：由于输入尺寸较大（1536），batch size设为2以避免OOM。

训练过程中可通过TensorBoard查看loss曲线和mAP变化：

tensorboard --logdir=runs/train

7. 推理与评估：测试模型性能

7.1 单图推理测试

使用训练好的模型进行单张图片检测：

python detect.py \ --source inference/images/test.jpg \ --weights runs/train/tph-yolov5-visdrone/weights/best.pt \ --img 1536 \ --conf 0.25 \ --name tph_result

结果将保存在runs/detect/tph_result目录下。

7.2 多尺度测试（MS-Testing）

为提升精度，启用多尺度测试策略：

python val.py \ --weights best.pt \ --data data/visdrone.yaml \ --img 1536 \ --task test \ --ms

其中--ms表示对同一图像缩放多个比例（1.3x, 1x, 0.83x, 0.67x）并水平翻转，共6次推理后融合结果。

7.3 模型集成（Ensemble）

训练多个略有差异的模型（不同输入尺寸、类别权重等），使用Weighted Boxes Fusion (WBF)融合预测框：

from utils.metrics import wbf_ensemble results = wbf_ensemble( weights=['m1.pt', 'm2.pt', 'm3.pt', 'm4.pt', 'm5.pt'], iou_thres=0.6, skip_box_thr=0.0001 )

相比NMS，WBF能保留更多有效框并加权合并，进一步提升mAP约0.5%-1%。