如何在PyTorch环境下微调ssd300_vgg16目标检测模型？

以下代码可直接在最新版PyTorch环境的Colab中运行，完全适配model=torchvision.models.detection.ssd300_vgg16(pretrained=False)的加载要求，支持自定义边界框、类别标签数据集训练。

1. 环境初始化与依赖校验 #

Colab默认预装匹配版本的PyTorch和Torchvision，首先运行代码校验GPU可用性：

import torch
import torchvision
from torchvision.models.detection import ssd300_vgg16
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import json
import os

# 设备配置
print(f"PyTorch版本: {torch.__version__}")
print(f"Torchvision版本: {torchvision.__version__}")
print(f"GPU可用状态: {torch.cuda.is_available()}")
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

注意：训练前先在Colab菜单「修改-笔记本设置」中把硬件加速器切换为GPU，否则训练速度极慢。

2. 自定义数据集类实现 #

Torchvision目标检测模型要求数据集返回格式为(图像张量, 标注字典)，标注字典必须包含边界框、类别标签字段，直接复用以下类即可：

class CustomDetDataset(Dataset):
    def __init__(self, data_root, anno_path):
        self.data_root = data_root
        # 加载标注文件
        with open(anno_path, 'r', encoding='utf-8') as f:
            self.anno_list = json.load(f)
        self.img_list = [item['filename'] for item in self.anno_list]

    def __getitem__(self, idx):
        # 读取图像并转张量
        img_path = os.path.join(self.data_root, 'images', self.img_list[idx])
        img = Image.open(img_path).convert('RGB')
        img_tensor = torchvision.transforms.functional.to_tensor(img)
        # 解析标注
        current_anno = self.anno_list[idx]
        boxes = torch.as_tensor(current_anno['bboxes'], dtype=torch.float32)
        labels = torch.as_tensor(current_anno['labels'], dtype=torch.int64)
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        iscrowd = torch.zeros((len(boxes),), dtype=torch.int64)
        image_id = torch.tensor([idx])

        target = {
            "boxes": boxes,
            "labels": labels,
            "image_id": image_id,
            "area": area,
            "iscrowd": iscrowd
        }
        return img_tensor, target

    def __len__(self):
        return len(self.img_list)

标注文件annotations.json格式参考，直接按这个结构整理自己的数据集即可：

[
  {
    "filename": "img_0001.jpg",
    "bboxes": [[12, 45, 230, 310], [89, 120, 450, 520]],
    "labels": [1, 2]
  }
]

注意：类别标签从1开始计数，0固定为背景类；边界框必须是[xmin, ymin, xmax, ymax]格式的原图绝对像素坐标，不能传归一化的0-1范围值，否则训练时loss会变为nan。

3. 模型加载与数据加载器配置 #

# 类别数配置：自定义类别总数 + 1（背景类），按自己的数据集修改
num_classes = 4  # 示例：3个自定义检测类+1个背景类
# 按要求加载无预训练权重的SSD300模型
model = ssd300_vgg16(pretrained=False, num_classes=num_classes)
model.to(device)

# 自定义batch拼接函数，适配不同图片的bbox数量不一致问题
def collate_fn(batch):
    return tuple(zip(*batch))

# 初始化数据集，替换为自己上传到Colab的数据集路径
dataset = CustomDetDataset(
    data_root='/content/your_dataset',
    anno_path='/content/your_dataset/annotations.json'
)
# 8:2切分训练、验证集
train_num = int(len(dataset)*0.8)
val_num = len(dataset) - train_num
train_set, val_set = torch.utils.data.random_split(dataset, [train_num, val_num])

train_loader = DataLoader(
    train_set,
    batch_size=8,  # T4显卡设8即可，V100可以开到16
    shuffle=True,
    num_workers=4,
    collate_fn=collate_fn
)
val_loader = DataLoader(
    val_set,
    batch_size=4,
    shuffle=False,
    num_workers=4,
    collate_fn=collate_fn
)

4. 训练循环实现 #

Torchvision封装的检测模型在train()模式下会自动计算损失，不需要手动实现损失函数：

# 优化器与学习率配置
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=5e-4)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

# 开始训练，从头训建议epoch设30以上，加载预训练权重微调设10即可
num_epochs = 15
for epoch in range(num_epochs):
    model.train()
    epoch_loss = 0
    for batch_idx, (images, targets) in enumerate(train_loader):
        # 数据移到GPU
        images = list(img.to(device) for img in images)
        targets = [{k:v.to(device) for k,v in t.items()} for t in targets]
        # 前向传播计算损失
        loss_dict = model(images, targets)
        total_loss = sum(loss for loss in loss_dict.values())
        # 反向传播更新参数
        optimizer.zero_grad()
        total_loss.backward()
        optimizer.step()

        epoch_loss += total_loss.item()
        if batch_idx % 10 == 0:
            print(f"Epoch:{epoch+1}/{num_epochs}  Batch:{batch_idx}  当前Loss:{total_loss.item():.4f}")
    lr_scheduler.step()
    print(f"Epoch {epoch+1} 结束，平均训练Loss:{epoch_loss/len(train_loader):.4f}")

# 保存训练好的权重
torch.save(model.state_dict(), '/content/ssd300_custom_trained.pth')

5. 推理示例 #

训练完成后可以用以下代码做单图预测：

model.eval()
# 加载测试图
test_img = Image.open('/content/test_img.jpg').convert('RGB')
test_tensor = torchvision.transforms.functional.to_tensor(test_img).to(device)
with torch.no_grad():
    pred = model([test_tensor])[0]
# 过滤置信度大于0.5的预测结果
keep = pred['scores'] > 0.5
res_boxes = pred['boxes'][keep].cpu().numpy()
res_labels = pred['labels'][keep].cpu().numpy()
res_scores = pred['scores'][keep].cpu().numpy()

关键注意事项 #

• 数据集如果比较大，不用每次重新上传，可以先挂载Google Drive，把路径改成Drive内的数据集路径即可
• 从头训练（pretrained=False）收敛速度很慢，对数据量要求高，如果没有特殊需求可以先设pretrained=True加载COCO预训练权重，再微调自定义类别，收敛速度和精度都会明显提升
• 训练过程如果出现显存不足，调小batch_size即可

内容的提问来源于stack exchange，提问作者user2140324