torchvision.models 是 torchvision 库中的核心模块之一，提供了大量预定义的计算机视觉模型，主要用于图像分类、目标检测、图像分割、视频分类等任务。这些模型包括经典的深度神经网络架构（如 ResNet、VGG）和现代模型（如 Vision Transformer），支持预训练权重，方便迁移学习和快速开发。所有模型都与 PyTorch 深度集成，继承自 torch.nn.Module，可以轻松嵌入到自定义的训练流程中。

以下是对 torchvision.models 模块的详尽介绍，涵盖其功能、支持的模型、子模块、用法、迁移学习、自定义修改、代码示例和注意事项。

---

1. torchvision.models 模块概述

torchvision.models 提供了开箱即用的深度学习模型，涵盖了计算机视觉领域的多种任务。它的主要特点包括：

• 预训练模型：许多模型提供预训练权重（通常在 ImageNet 数据集上预训练），适合迁移学习或特征提取。
• 模块化设计：模型结构清晰，支持自定义修改（如替换分类层）。
• 多任务支持：不仅限于图像分类，还包括目标检测、图像分割、视频分类等。
• 易于使用：通过简单的函数调用即可加载模型，并与 PyTorch 的其他组件（如优化器、数据集）无缝集成。
• 轻量级模型：提供量化模型（如 MobileNet），适合资源受限的设备。

---

2. torchvision.models 的子模块

torchvision.models 按任务类型组织为以下几个子模块：

(1) 图像分类 (torchvision.models)

提供经典和现代的图像分类模型，适用于从简单到复杂的分类任务。

• 常用模型：

  • AlexNet：早期的深度卷积网络，2012 年 ImageNet 冠军。
  • VGG（VGG11, VGG16, VGG19 等）：深层卷积网络，强调深度和一致的卷积核。
  • ResNet（ResNet18, ResNet50, ResNet101 等）：引入残差连接，解决深层网络的退化问题。
  • Inception V3：通过多尺度卷积提高效率。
  • DenseNet：密集连接网络，减少参数量并增强特征复用。
  • MobileNet（MobileNetV2, MobileNetV3）：轻量级模型，适合移动设备。
  • EfficientNet：通过复合缩放优化模型性能。
  • Vision Transformer (ViT)：基于 Transformer 的图像分类模型，适合大规模数据。
  • Swin Transformer：分层 Transformer 架构，性能优异。

• 特点：

  • 所有模型支持 pretrained=True（或 weights='IMAGENET1K_V1'）加载 ImageNet 预训练权重。
  • 模型输出通常为 1000 类（ImageNet 类别数），需要替换分类头以适应自定义任务。

• 使用示例：

  import torchvision.models as models
  
  # 加载预训练的 ResNet50
  model = models.resnet50(weights='IMAGENET1K_V1')  # 或 pretrained=True（旧版）
  model.eval()  # 设置为评估模式
  
  # 检查模型结构
  print(model)
  

(2) 目标检测 (torchvision.models.detection)

提供用于目标检测和实例分割的模型，支持边界框预测和掩码生成。

• 常用模型：

  • Faster R-CNN：基于 ResNet-50 的目标检测模型，支持边界框预测。
  • Mask R-CNN：在 Faster R-CNN 基础上增加实例分割功能。
  • Keypoint R-CNN：支持人体关键点检测。
  • RetinaNet：单阶段检测器，基于 Focal Loss 解决前景-背景不平衡问题。
  • SSD (Single Shot MultiBox Detector)：轻量级单阶段检测模型。

• 特点：

  • 支持 COCO 数据集的预训练权重。
  • 输入为图像列表，输出包含边界框、标签和置信度（以及掩码，视模型而定）。
  • 需要特定的数据预处理（如归一化）和后处理（如 NMS）。

• 使用示例：

  from torchvision.models.detection import fasterrcnn_resnet50_fpn
  
  # 加载预训练的 Faster R-CNN
  model = fasterrcnn_resnet50_fpn(weights='COCO_V1')
  model.eval()
  
  # 输入示例（假设 images 是张量列表）
  import torch
  images = [torch.rand(3, 800, 800)]
  predictions = model(images)  # 输出包含 boxes、labels、scores
  

(3) 图像分割 (torchvision.models.segmentation)

提供语义分割和实例分割模型，输出像素级预测。

• 常用模型：

  • FCN (Fully Convolutional Network)：基于 ResNet-50 或 ResNet-101，支持语义分割。
  • DeepLabV3：使用空洞卷积提高感受野，性能更强。
  • LR-ASPP：轻量级分割模型，适合资源受限场景。

• 特点：

  • 支持 COCO 或 VOC 数据集的预训练权重。
  • 输出为每个像素的类别预测（语义分割）或掩码（实例分割）。
  • 需要调整输入图像大小以适应模型要求。

• 使用示例：

  from torchvision.models.segmentation import deeplabv3_resnet50
  
  # 加载预训练的 DeepLabV3
  model = deeplabv3_resnet50(weights='COCO_WITH_VOC_LABELS_V1')
  model.eval()
  
  # 输入示例
  image = torch.rand(1, 3, 520, 520)
  output = model(image)['out']  # 输出形状为 [batch, num_classes, H, W]
  

(4) 视频分类 (torchvision.models.video)

提供用于视频分类的模型，处理时序数据。

• 常用模型：

  • R3D：3D ResNet，基于 3D 卷积处理视频帧。
  • MC3：混合 3D/2D 卷积模型。
  • R2Plus1D：分解 3D 卷积为 2D+1D，减少计算量。

• 特点：

  • 支持 Kinetics 数据集的预训练权重。
  • 输入为视频帧序列（形状为 [batch, channels, time, height, width]）。

• 使用示例：

  from torchvision.models.video import r3d_18
  
  # 加载预训练的 R3D-18
  model = r3d_18(weights='KINETICS400_V1')
  model.eval()
  
  # 输入示例（假设 video 是视频张量）
  video = torch.rand(1, 3, 16, 112, 112)  # 16 帧
  output = model(video)
  

(5) 量化模型 (torchvision.models.quantization)

提供量化的轻量级模型，适合在边缘设备上部署。

• 常用模型：

  • Quantized ResNet、Quantized MobileNetV2 等。
  • 使用 int8 量化减少模型大小和推理时间。

• 使用示例：

  from torchvision.models.quantization import mobilenet_v2
  
  # 加载量化模型
  model = mobilenet_v2(weights='IMAGENET1K_V1', quantize=True)
  model.eval()
  

---

3. 加载模型的通用参数

大多数 torchvision.models 函数支持以下参数：

• weights (str or None)：指定预训练权重（如 ‘IMAGENET1K_V1’、‘COCO_V1’），若为 None 则随机初始化。
  • 旧版使用 pretrained=True，新版推荐使用 weights。
• num_classes (int)：自定义分类模型的输出类别数（默认 1000，适用于 ImageNet）。
• progress (bool)：是否显示权重下载进度（默认 True）。
• 其他参数：某些模型支持特定参数，如 aux_loss（Inception V3）、box_score_thresh（检测模型）等。

---

4. 迁移学习与模型修改

torchvision.models 的预训练模型非常适合迁移学习。以下是两种常见用法：

(1) 特征提取

冻结预训练模型的卷积层，仅训练分类头：

import torchvision.models as models
import torch.nn as nn

# 加载 ResNet50
model = models.resnet50(weights='IMAGENET1K_V1')

# 冻结所有参数
for param in model.parameters():
    param.requires_grad = False

# 替换分类头
num_classes = 10  # 假设新任务有 10 个类别
model.fc = nn.Linear(model.fc.in_features, num_classes)

# 仅训练分类头
optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.01, momentum=0.9)

(2) 微调 (Fine-tuning)

对整个模型进行微调，通常使用较小的学习率：

model = models.resnet50(weights='IMAGENET1K_V1')
model.fc = nn.Linear(model.fc.in_features, num_classes)

# 使用小的学习率微调所有参数
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

(3) 自定义模型结构

可以修改模型的某些层，例如添加 dropout 或替换骨干网络：

model = models.resnet50(weights='IMAGENET1K_V1')
model.fc = nn.Sequential(
    nn.Linear(model.fc.in_features, 512),
    nn.ReLU(),
    nn.Dropout(0.5),
    nn.Linear(512, num_classes)
)

---

5. 结合 DataLoader 和训练

以下是一个完整的图像分类训练示例：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, models, transforms
from torch.utils.data import DataLoader

# 数据预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载数据集
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# 加载和修改模型
model = models.resnet18(weights='IMAGENET1K_V1')
model.fc = nn.Linear(model.fc.in_features, 10)  # CIFAR-10 有 10 个类别

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 训练循环
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
model.train()

for epoch in range(5):  # 训练 5 个 epoch
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
    print(f'Epoch {epoch+1}, Loss: {loss.item()}')

---

6. 注意事项

• 版本兼容性：确保 torchvision 和 PyTorch 版本匹配（截至 2025 年 10 月，推荐 torchvision 0.18.0 和 PyTorch 2.3.0）。
• 权重加载：
  • 新版使用 weights='IMAGENET1K_V1' 而非 pretrained=True。
  • 部分模型（如检测和分割模型）可能需要特定数据集的权重（如 COCO）。
• 输入要求：
  • 分类模型通常要求输入为 [batch, 3, H, W]，H 和 W 至少为 224（某些模型如 ViT 要求特定大小）。
  • 检测模型接受图像列表，需手动归一化。
• 内存与性能：
  • 大型模型（如 ResNet101、DeepLabV3）需要较多 GPU 内存。
  • 量化模型适合边缘设备，但精度可能略有下降。
• 后处理：
  • 检测和分割模型的输出需要后处理（如 NMS、非极大值抑制）来过滤冗余预测。

---

7. 学习建议

• 初学者：从 ResNet 或 MobileNet 开始，学习如何加载预训练模型和进行迁移学习。
• 进阶：尝试 Faster R-CNN 或 Mask R-CNN，结合 COCO 数据集进行目标检测或分割。
• 高级：探索 Vision Transformer 或 Swin Transformer，学习如何处理大规模数据和复杂任务。
• 实践：结合 torchvision.datasets 和 transforms，构建完整的训练 pipeline。

---

8. 总结

torchvision.models 模块提供了丰富的预训练模型，覆盖图像分类、目标检测、图像分割和视频分类等多种任务。通过支持预训练权重和灵活的修改接口，它非常适合快速原型开发和迁移学习。无论是初学者还是高级开发者，都可以通过该模块快速构建高效的计算机视觉模型。