Swin Transformer目标检测实战：从环境配置到模型训练（CUDA10.2+Python3.6保姆级教程）
·
Swin Transformer目标检测实战：从零构建高效识别系统
最近在做一个工业质检的项目，客户要求对微小缺陷进行高精度识别，传统的CNN模型在长距离依赖建模上总是差那么点意思。尝试了几个方案后，我把目光投向了Swin Transformer——这个在图像分类任务上大放异彩的视觉Transformer变体。但说实话，刚开始把它应用到目标检测时，踩的坑可真不少：环境配置的版本冲突、配置文件的理解偏差、训练过程中的内存溢出……折腾了两周才真正跑通整个流程。
如果你也正在考虑将Swin Transformer用于目标检测任务，无论是学术研究还是工业落地，这篇文章或许能帮你省下不少时间。我不会简单重复官方文档的内容，而是结合自己实际项目中的经验，从环境搭建、代码修改到训练调优，一步步带你构建一个可用的Swin Transformer目标检测系统。我们重点会放在那些容易出错的细节上，特别是当你使用较老但稳定的CUDA 10.2环境时需要注意的事项。
1. 环境搭建：避开版本依赖的“雷区”
很多教程会告诉你“按顺序安装这些包就行”，但实际过程中，版本兼容性问题往往是第一个拦路虎。特别是PyTorch、CUDA和mmcv-full这三者的匹配，稍有偏差就会导致后续训练失败。
1.1 虚拟环境与PyTorch选择
我强烈建议使用Anaconda创建独立的虚拟环境，这能避免与系统已有Python包产生冲突。虽然Python 3.6-3.8都可以，但我个人更倾向于3.7，它在兼容性和新特性之间取得了不错的平衡。
conda create -n swin_det python=3.7 -y
conda activate swin_det
接下来是PyTorch安装。这里有个关键点：不要盲目追求最新版本。对于CUDA 10.2，PyTorch 1.8.0是个比较稳妥的选择，它既有较好的性能，又有相对完善的生态支持。
pip install torch==1.8.0 torchvision==0.9.0
如果你发现安装后无法调用GPU，可以验证一下：
import torch
print(torch.__version__) # 应该显示1.8.0
print(torch.cuda.is_available()) # 应该返回True
注意：有些服务器预装了不同版本的CUDA，你可以通过
nvcc --version查看CUDA版本。如果显示的是11.x，那么需要对应安装PyTorch 1.9.0或更高版本。
1.2 mmcv-full的“正确”安装方式
mmcv-full是OpenMMLab系列工具包的核心，但它的安装方式让很多新手困惑。实际上，根据PyTorch和CUDA版本选择正确的安装命令至关重要。
对于PyTorch 1.8.0 + CUDA 10.2的组合，我推荐使用以下命令：
pip install mmcv-full==1.7.0 -f https://download.openmmlab.com/mmcv/dist/cu102/torch1.8.0/index.html
这个命令明确指定了CUDA版本（cu102）和PyTorch版本（torch1.8.0），确保下载预编译的、兼容的二进制包。如果网络环境不佳，也可以尝试使用openmim：
pip install openmim
mim install mmcv-full==1.7.0
安装完成后，可以简单测试：
import mmcv
print(mmcv.__version__) # 应该显示1.7.0
1.3 其他依赖的安装技巧
pycocotools是COCO数据集评估必需的库，但在某些系统上直接pip install可能会失败。我的经验是先安装编译依赖：
# Ubuntu/Debian系统
sudo apt-get install cython
# CentOS/RHEL系统
sudo yum install cython
# 然后从源码安装
git clone https://github.com/cocodataset/cocoapi
cd cocoapi/PythonAPI
python setup.py build_ext install
apex库用于混合精度训练，能显著减少显存占用并加速训练。但它的安装需要特别注意版本匹配：
git clone https://github.com/NVIDIA/apex
cd apex
# 对于PyTorch 1.8.0，使用最新版本即可
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
如果编译失败，可能是CUDA版本或编译器问题。可以尝试不编译CUDA扩展的简化安装：
pip install -v --no-cache-dir ./
最后安装MMDetection：
git clone https://github.com/open-mmlab/mmdetection.git
cd mmdetection
pip install -r requirements/build.txt
pip install -v -e .
提示：
-e参数表示以“可编辑”模式安装，这样你修改源码后无需重新安装就能生效。
2. 数据集准备与配置文件解析
环境搭好了，接下来要处理数据和配置文件。这部分看似简单，但配置错误会导致训练无法开始或结果异常。
2.1 自定义数据集的正确格式
假设我们有一个工业缺陷检测数据集，包含“划痕”、“凹陷”、“污渍”三类。我们需要将数据转换为COCO格式，这是MMDetection支持的标准格式之一。
COCO格式的核心是annotations字段，每个标注需要包含以下信息：
{
"images": [
{
"id": 1,
"file_name": "defect_001.jpg",
"height": 800,
"width": 1200
}
],
"annotations": [
{
"id": 1,
"image_id": 1,
"category_id": 0,
"bbox": [x, y, width, height], // [左上角x, 左上角y, 宽度, 高度]
"area": width * height,
"iscrowd": 0
}
],
"categories": [
{"id": 0, "name": "scratch"},
{"id": 1, "name": "dent"},
{"id": 2, "name": "stain"}
]
}
实际项目中，我通常使用以下Python脚本进行格式转换：
import json
from pathlib import Path
def convert_to_coco(annotations, output_path):
coco_format = {
"images": [],
"annotations": [],
"categories": []
}
# 添加类别信息
categories = ["scratch", "dent", "stain"]
for i, cat_name in enumerate(categories):
coco_format["categories"].append({
"id": i,
"name": cat_name,
"supercategory": "defect"
})
# 转换每个图像的标注
ann_id = 1
for img_id, anns in enumerate(annotations):
# 添加图像信息
coco_format["images"].append({
"id": img_id,
"file_name": anns["filename"],
"height": anns["height"],
"width": anns["width"]
})
# 添加该图像的所有标注
for ann in anns["bboxes"]:
x, y, w, h = ann["bbox"]
coco_format["annotations"].append({
"id": ann_id,
"image_id": img_id,
"category_id": ann["category_id"],
"bbox": [x, y, w, h],
"area": w * h,
"iscrowd": 0
})
ann_id += 1
with open(output_path, 'w') as f:
json.dump(coco_format, f, indent=2)
2.2 配置文件的多层结构理解
MMDetection的配置文件采用继承机制，理解这个结构能帮你快速定位需要修改的地方。以Swin-Tiny为例，配置文件通常是这样组织的：
configs/swin/
├── mask_rcnn_swin_tiny_patch4_window7_mstrain_480-800_adamw_1x_coco.py
├── _base_/
│ ├── datasets/coco_detection.py
│ ├── models/mask_rcnn_swin_fpn.py
│ ├── schedules/schedule_1x.py
│ └── default_runtime.py
这种继承关系意味着：
- 主配置文件继承_base_中的各个组件
- 修改_base_中的文件会影响所有使用该组件的配置
- 你可以在主配置文件中覆盖任何父配置
对于纯目标检测任务（不需要实例分割），我们需要修改几个关键位置。首先找到models/mask_rcnn_swin_fpn.py，注释掉与mask相关的部分：
# 注释掉以下mask相关配置
# mask_roi_extractor=dict(
# type='SingleRoIExtractor',
# roi_layer=dict(type='RoIAlign', output_size=14, sampling_ratio=0),
# out_channels=256,
# featmap_strides=[4, 8, 16, 32]),
# mask_head=dict(
# type='FCNMaskHead',
# num_convs=4,
# in_channels=256,
# conv_out_channels=256,
# num_classes=80,
# loss_mask=dict(
# type='CrossEntropyLoss', use_mask=True, loss_weight=1.0))
注意：注释时要确保括号匹配，否则会导致配置文件解析失败。建议使用代码编辑器的括号高亮功能辅助检查。
2.3 关键参数调整策略
不同的数据集和应用场景需要不同的参数设置。以下是一些经验值：
| 参数 | 小数据集 (<1k图像) | 中等数据集 (1k-10k) | 大数据集 (>10k) |
|---|---|---|---|
| batch_size | 2-4 | 4-8 | 8-16 |
| learning_rate | 0.0001 | 0.0002 | 0.0005 |
| warmup_iters | 500 | 1000 | 2000 |
| max_epochs | 12-24 | 24-36 | 36-48 |
| optimizer | AdamW | AdamW | SGD |
对于类别数的修改，需要同步更新两个地方：
- 在
models/mask_rcnn_swin_fpn.py中（修改后应重命名为swin_fpn.py以避免混淆）：
# 第54行附近
roi_head=dict(
type='StandardRoIHead',
bbox_roi_extractor=...,
bbox_head=dict(
type='Shared2FCBBoxHead',
num_classes=3, # 改为你的类别数
...),
# mask_head部分已注释掉
)
- 在
datasets/coco.py中修改类别名称：
CLASSES = ('scratch', 'dent', 'stain') # 你的类别名称
3. 模型训练：从单卡到多卡
配置完成后，终于可以开始训练了。但训练过程中也有很多技巧和注意事项。
3.1 单卡训练与调试
对于初次尝试或调试，建议从单卡开始：
python tools/train.py configs/swin/mask_rcnn_swin_tiny_patch4_window7_mstrain_480-800_adamw_1x_coco.py \
--work-dir work_dirs/swin_tiny_det \
--cfg-options data.samples_per_gpu=2 \
data.workers_per_gpu=2 \
runner.max_epochs=24
这里有几个实用参数：
--work-dir: 指定输出目录，包含日志、检查点等--cfg-options: 临时覆盖配置参数，无需修改配置文件--resume-from: 从某个检查点恢复训练--no-validate: 训练期间不进行验证（加速训练）
训练开始后，关注几个关键指标：
# 在训练日志中你会看到类似信息
# 每50个iteration输出一次
[>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>] 100/100, 2.0 task/s, elapsed: 50s, ETA: 0s
Epoch [1][50/100] lr: 1.00e-04, eta: 1:23:45, time: 0.512, data_time: 0.012, memory: 5123, loss_rpn_cls: 0.1234, loss_rpn_bbox: 0.0567, loss_cls: 0.2345, loss_bbox: 0.0890, loss: 0.5036
- loss: 总损失，应该随着训练逐渐下降
- loss_cls: 分类损失，反映模型识别类别的能力
- loss_bbox: 边界框回归损失，反映定位精度
- memory: 显存占用（MB），用于监控是否溢出
3.2 多卡分布式训练
当单卡训练稳定后，可以扩展到多卡加速。MMDetection使用dist_train.sh脚本：
./tools/dist_train.sh configs/swin/mask_rcnn_swin_tiny_patch4_window7_mstrain_480-800_adamw_1x_coco.py 4 \
--work-dir work_dirs/swin_tiny_det_4gpu \
--cfg-options data.samples_per_gpu=4 \
data.workers_per_gpu=4 \
optimizer.lr=0.0008
这里有个重要公式：学习率需要随GPU数量线性缩放。如果单卡学习率是0.0002，那么4卡时应设为0.0008。
分布式训练常见问题及解决方案：
-
NCCL错误：通常是网络问题，尝试设置
export NCCL_IB_DISABLE=1 export NCCL_SOCKET_IFNAME=eth0 # 根据实际网卡名调整 -
显存不均：某些GPU显存占用过高，可以尝试
export CUDA_VISIBLE_DEVICES=0,1,2,3 # 明确指定GPU -
训练速度不线性：数据加载可能成为瓶颈，增加
workers_per_gpu或使用更快的存储
3.3 训练监控与调优
训练过程中，我习惯使用TensorBoard监控训练状态：
tensorboard --logdir work_dirs/swin_tiny_det --port 6006
然后在浏览器打开localhost:6006，可以看到：
- Loss曲线：观察是否收敛，是否有震荡
- 学习率曲线：检查warmup和衰减策略是否正确执行
- 验证集mAP：最重要的指标，反映模型泛化能力
如果发现过拟合（训练loss下降但验证集指标不升反降），可以尝试：
# 在配置文件中增加正则化
model = dict(
...
train_cfg=dict(
rpn=dict(
assigner=dict(
type='MaxIoUAssigner',
pos_iou_thr=0.7,
neg_iou_thr=0.3,
min_pos_iou=0.3,
match_low_quality=True,
ignore_iof_thr=-1),
sampler=dict(
type='RandomSampler',
num=256,
pos_fraction=0.5,
neg_pos_ub=-1,
add_gt_as_proposals=False),
allowed_border=-1,
pos_weight=-1,
debug=False),
rpn_proposal=dict(
nms_pre=2000,
max_per_img=1000,
nms=dict(type='nms', iou_threshold=0.7),
min_bbox_size=0),
rcnn=dict(
assigner=dict(
type='MaxIoUAssigner',
pos_iou_thr=0.5,
neg_iou_thr=0.5,
min_pos_iou=0.5,
match_low_quality=False,
ignore_iof_thr=-1),
sampler=dict(
type='RandomSampler',
num=512,
pos_fraction=0.25,
neg_pos_ub=-1,
add_gt_as_proposals=True),
pos_weight=-1,
debug=False)),
test_cfg=dict(
rpn=dict(
nms_pre=1000,
max_per_img=1000,
nms=dict(type='nms', iou_threshold=0.7),
min_bbox_size=0),
rcnn=dict(
score_thr=0.05,
nms=dict(type='nms', iou_threshold=0.5),
max_per_img=100)
)
)
4. 推理部署与性能优化
模型训练完成后，我们需要在实际场景中使用它。这部分涉及模型测试、性能评估和优化。
4.1 单张图像推理
MMDetection提供了方便的demo脚本，但实际项目中我们通常需要自定义推理流程：
import mmcv
from mmdet.apis import init_detector, inference_detector
import numpy as np
def detect_single_image(config_file, checkpoint_file, image_path, score_thr=0.3):
# 初始化模型
model = init_detector(config_file, checkpoint_file, device='cuda:0')
# 读取图像
img = mmcv.imread(image_path)
# 推理
result = inference_detector(model, img)
# 解析结果
bboxes = []
labels = []
scores = []
for class_id, class_result in enumerate(result):
if len(class_result) > 0:
for bbox in class_result:
if bbox[4] >= score_thr: # 置信度阈值过滤
bboxes.append(bbox[:4].astype(int))
labels.append(class_id)
scores.append(bbox[4])
# 非极大值抑制（NMS）
if len(bboxes) > 0:
keep_indices = mmcv.ops.nms(
np.array(bboxes).astype(np.float32),
np.array(scores),
iou_threshold=0.5
)
bboxes = [bboxes[i] for i in keep_indices]
labels = [labels[i] for i in keep_indices]
scores = [scores[i] for i in keep_indices]
return bboxes, labels, scores
# 使用示例
bboxes, labels, scores = detect_single_image(
'configs/swin/mask_rcnn_swin_tiny_patch4_window7_mstrain_480-800_adamw_1x_coco.py',
'work_dirs/swin_tiny_det/latest.pth',
'test_image.jpg',
score_thr=0.5
)
对于视频推理，可以这样处理：
def detect_video(config_file, checkpoint_file, video_path, output_path):
model = init_detector(config_file, checkpoint_file, device='cuda:0')
video = mmcv.VideoReader(video_path)
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
fps = video.fps
frame_size = (video.width, video.height)
out_video = cv2.VideoWriter(output_path, fourcc, fps, frame_size)
for frame in mmcv.track_iter_progress(video):
result = inference_detector(model, frame)
vis_frame = model.show_result(frame, result, score_thr=0.5)
out_video.write(vis_frame)
out_video.release()
4.2 批量测试与评估
对于整个测试集的评估，使用官方脚本：
# 单卡测试
python tools/test.py configs/swin/mask_rcnn_swin_tiny_patch4_window7_mstrain_480-800_adamw_1x_coco.py \
work_dirs/swin_tiny_det/latest.pth \
--eval bbox \
--out results.pkl \
--show-dir results_vis
# 多卡测试
./tools/dist_test.sh configs/swin/mask_rcnn_swin_tiny_patch4_window7_mstrain_480-800_adamw_1x_coco.py \
work_dirs/swin_tiny_det/latest.pth \
4 \
--eval bbox \
--eval-options "classwise=True"
评估结果会显示多个指标：
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.425
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.687
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.456
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.234
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.467
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.589
关键指标解读：
- mAP@[0.5:0.95]：主要评价指标，IoU阈值从0.5到0.95的平均精度
- mAP@0.5：宽松指标，IoU阈值0.5时的精度
- mAP@0.75：严格指标，IoU阈值0.75时的精度
- 小/中/大目标AP：反映模型对不同尺度目标的检测能力
4.3 模型优化与加速
在实际部署中，我们往往需要权衡精度和速度。以下是一些优化策略：
1. 模型剪枝与量化
import torch
from mmdet.apis import init_detector
# 加载训练好的模型
model = init_detector(config_file, checkpoint_file, device='cpu')
# 动态量化（PyTorch内置）
quantized_model = torch.quantization.quantize_dynamic(
model,
{torch.nn.Linear, torch.nn.Conv2d},
dtype=torch.qint8
)
# 保存量化模型
torch.save(quantized_model.state_dict(), 'quantized_model.pth')
2. TensorRT加速
对于生产环境，TensorRT能提供显著的推理加速。首先将模型转换为ONNX格式：
from mmdet.apis import init_detector
import torch
model = init_detector(config_file, checkpoint_file, device='cuda')
input_shape = (1, 3, 800, 1333) # 根据你的输入尺寸调整
# 导出ONNX
torch.onnx.export(
model,
torch.randn(input_shape).cuda(),
'model.onnx',
input_names=['input'],
output_names=['output'],
dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
)
然后使用TensorRT转换：
# 安装TensorRT
# 参考NVIDIA官方文档：https://developer.nvidia.com/tensorrt
# 转换ONNX到TensorRT
trtexec --onnx=model.onnx \
--saveEngine=model.trt \
--fp16 \
--workspace=4096
3. 推理优化技巧
- 批处理：合理设置batch_size，充分利用GPU并行能力
- 异步推理：使用多线程实现数据加载、推理、后处理的流水线
- 内存复用：避免频繁的内存分配和释放
import threading
import queue
import torch
class AsyncInference:
def __init__(self, model, batch_size=4, num_workers=2):
self.model = model
self.batch_size = batch_size
self.input_queue = queue.Queue(maxsize=10)
self.output_queue = queue.Queue(maxsize=10)
self.workers = []
for _ in range(num_workers):
worker = threading.Thread(target=self._worker)
worker.daemon = True
worker.start()
self.workers.append(worker)
def _worker(self):
while True:
batch_data = self.input_queue.get()
if batch_data is None:
break
with torch.no_grad():
results = self.model(batch_data)
self.output_queue.put(results)
def put(self, data):
self.input_queue.put(data)
def get(self):
return self.output_queue.get()
5. 实际项目中的问题排查
即使按照教程一步步操作，在实际项目中还是会遇到各种问题。这里分享一些常见问题的排查经验。
5.1 训练过程中的常见错误
问题1：CUDA out of memory
这是最常见的问题。解决方案：
- 减小
batch_size（samples_per_gpu） - 减小输入图像尺寸
- 使用梯度累积：
# 在配置文件中 optimizer_config = dict( type='OptimizerHook', grad_clip=None, cumulative_iters=4 # 每4次迭代更新一次权重 ) - 使用混合精度训练（需要apex库）：
# 训练命令中添加 --fp16
问题2：Loss为NaN或异常大
可能原因和解决方案：
- 学习率过大：尝试减小10倍
- 数据中有异常值：检查标注数据
- 梯度爆炸：添加梯度裁剪
optimizer_config = dict( type='OptimizerHook', grad_clip=dict(max_norm=35, norm_type=2) )
问题3：验证集指标不提升
- 过拟合：增加数据增强，添加Dropout
- 学习率策略不当：尝试CosineAnnealing或ReduceLROnPlateau
- 模型容量不足：尝试更大的Swin模型（如Swin-Small、Swin-Base）
5.2 推理时的性能问题
问题：推理速度慢
优化策略：
- 使用更小的模型（Swin-Tiny vs Swin-Base）
- 减小输入分辨率
- 使用TensorRT或ONNX Runtime加速
- 优化后处理（NMS）代码：
# 自定义更快的NMS实现
def fast_nms(boxes, scores, iou_threshold=0.5):
"""基于PyTorch的快速NMS实现"""
if len(boxes) == 0:
return torch.empty((0,), dtype=torch.long)
# 按分数降序排序
_, indices = scores.sort(descending=True)
boxes = boxes[indices]
keep = []
while boxes.size(0) > 0:
keep.append(indices[0].item())
if boxes.size(0) == 1:
break
# 计算IoU
ious = bbox_iou(boxes[0:1], boxes[1:])
# 保留IoU小于阈值的框
mask = ious[0] < iou_threshold
boxes = boxes[1:][mask]
indices = indices[1:][mask]
return torch.tensor(keep, dtype=torch.long)
5.3 模型精度调优技巧
如果模型精度达不到预期，可以尝试：
-
数据增强策略调整：
# 在配置文件中增强数据增强 train_pipeline = [ dict(type='LoadImageFromFile'), dict(type='LoadAnnotations', with_bbox=True), dict(type='Resize', img_scale=(1333, 800), keep_ratio=True), dict(type='RandomFlip', flip_ratio=0.5), dict(type='RandomCrop', crop_size=(0.8, 0.8)), # 随机裁剪 dict(type='PhotoMetricDistortion', # 光度失真 brightness_delta=32, contrast_range=(0.5, 1.5), saturation_range=(0.5, 1.5), hue_delta=18), dict(type='Normalize', mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375]), dict(type='Pad', size_divisor=32), dict(type='DefaultFormatBundle'), dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels']) ] -
损失函数调整：
# 使用Focal Loss处理类别不平衡 bbox_head=dict( type='Shared2FCBBoxHead', num_classes=3, reg_decoded_bbox=True, loss_cls=dict( type='FocalLoss', use_sigmoid=True, gamma=2.0, alpha=0.25, loss_weight=1.0), loss_bbox=dict(type='GIoULoss', loss_weight=1.3)) -
多尺度训练与测试：
# 多尺度训练 img_scale=[(1333, 800), (1600, 960), (1920, 1152)], # 多尺度测试（TTA） test_pipeline = [ dict(type='LoadImageFromFile'), dict( type='MultiScaleFlipAug', img_scale=[(1333, 800), (1600, 960), (1920, 1152)], flip=True, transforms=[ dict(type='Resize', keep_ratio=True), dict(type='RandomFlip'), dict(type='Normalize', mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375]), dict(type='Pad', size_divisor=32), dict(type='ImageToTensor', keys=['img']), dict(type='Collect', keys=['img']) ]) ]
在实际的工业质检项目中，我最终选择Swin-Small作为基础模型，输入分辨率保持1333×800，使用多尺度训练和Focal Loss。经过24个epoch的训练，在验证集上达到了0.68的mAP，推理速度在V100上达到23FPS，完全满足了产线实时检测的需求。最大的收获不是调出了多高的精度，而是理解了整个流程中每个环节的权衡——有时候为了5%的精度提升，推理速度可能会下降50%，这时候就需要根据实际场景做出选择了。
DAMO开发者矩阵，由阿里巴巴达摩院和中国互联网协会联合发起，致力于探讨最前沿的技术趋势与应用成果，搭建高质量的交流与分享平台，推动技术创新与产业应用链接，围绕“人工智能与新型计算”构建开放共享的开发者生态。
更多推荐
4
0- 0
已为社区贡献2条内容
所有评论(0)