Skip to content
←

필사 모드: torchvision 完全指南 — 从图像分类到 Object Detection、Segmentation

中文
0%
정확도 0%
💡 왼쪽 원문을 읽으면서 오른쪽에 따라 써보세요. Tab 키로 힌트를 받을 수 있습니다.
torchvision Guide

引言

torchvision 是 PyTorch 官方的计算机视觉库。图像变换、预训练模型、数据集，直到 Object Detection — CV 所需的几乎一切都包含在内。

pip install torch torchvision

Part 1: Transforms v2 — 图像预处理的革新

基本变换

import torch
from torchvision import transforms
from torchvision.transforms import v2  # 推荐使用 v2！
from PIL import Image

# Transforms v2（最新、推荐）
transform = v2.Compose([
    v2.RandomResizedCrop(224),         # 随机裁剪 + 缩放
    v2.RandomHorizontalFlip(p=0.5),    # 50% 概率左右翻转
    v2.ColorJitter(                     # 颜色扰动
        brightness=0.2,
        contrast=0.2,
        saturation=0.2,
        hue=0.1
    ),
    v2.ToImage(),                       # PIL → Tensor 图像
    v2.ToDtype(torch.float32, scale=True),  # [0, 255] → [0.0, 1.0]
    v2.Normalize(                       # ImageNet 归一化
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    ),
])

img = Image.open("cat.jpg")
tensor = transform(img)
print(tensor.shape)  # torch.Size([3, 224, 224])

v2 的核心 — Bounding Box + Mask 同时变换

# v1 中只变换图像 → 导致 BBox 错位！
# v2 中图像 + BBox + Mask + Label 同时变换

from torchvision import tv_tensors

# Object Detection 用变换
det_transform = v2.Compose([
    v2.RandomHorizontalFlip(p=0.5),
    v2.RandomPhotometricDistort(),
    v2.RandomZoomOut(fill={tv_tensors.Image: (123, 117, 104)}),
    v2.RandomIoUCrop(),
    v2.SanitizeBoundingBoxes(),  # 去除无效 bbox
    v2.ToDtype(torch.float32, scale=True),
])

# 把图像 + bbox 一起变换时，bbox 也会自动跟随！
image = tv_tensors.Image(torch.randint(0, 256, (3, 500, 500), dtype=torch.uint8))
boxes = tv_tensors.BoundingBoxes(
    [[100, 100, 300, 300], [200, 200, 400, 400]],
    format="XYXY",
    canvas_size=(500, 500)
)
labels = torch.tensor([1, 2])

# 同时变换！
out_img, out_boxes, out_labels = det_transform(image, boxes, labels)

常用的 Augmentation 配方

# 训练用（强 augmentation）
train_transform = v2.Compose([
    v2.RandomResizedCrop(224, scale=(0.6, 1.0)),
    v2.RandomHorizontalFlip(),
    v2.RandAugment(num_ops=2, magnitude=9),  # AutoAugment 系列
    v2.ToImage(),
    v2.ToDtype(torch.float32, scale=True),
    v2.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    v2.RandomErasing(p=0.25),  # CutOut 效果
])

# 验证/推理用（无变形）
val_transform = v2.Compose([
    v2.Resize(256),
    v2.CenterCrop(224),
    v2.ToImage(),
    v2.ToDtype(torch.float32, scale=True),
    v2.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

Part 2: 预训练模型（Model Zoo）

图像分类模型

from torchvision import models
from torchvision.models import (
    ResNet50_Weights, EfficientNet_V2_S_Weights,
    ViT_B_16_Weights, ConvNeXt_Small_Weights
)

# ResNet-50 (2015, CNN 的基本功)
resnet = models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)
# Top-1: 80.858%, 参数量: 25.6M

# EfficientNet V2 (2021, 效率之王)
effnet = models.efficientnet_v2_s(weights=EfficientNet_V2_S_Weights.IMAGENET1K_V1)
# Top-1: 84.228%, 参数量: 21.5M

# Vision Transformer (2020, Transformer 进军 CV)
vit = models.vit_b_16(weights=ViT_B_16_Weights.IMAGENET1K_SWAG_E2E_V1)
# Top-1: 85.304%, 参数量: 86.6M

# ConvNeXt (2022, CNN 的反击 — 把 Transformer 的手法用到 CNN 上)
convnext = models.convnext_small(weights=ConvNeXt_Small_Weights.IMAGENET1K_V1)
# Top-1: 83.616%, 参数量: 50.2M

模型选择指南:
├── 需要快速推理 → MobileNet V3 / EfficientNet-Lite
├── 精度优先 → ViT-L / Swin Transformer V2
├── 均衡（实务推荐） → EfficientNet V2 / ConvNeXt
└── 学习/理解目的 → ResNet-50（基本功）

推理（Inference）

from torchvision.models import ViT_B_16_Weights

# 加载模型 + 预处理
weights = ViT_B_16_Weights.IMAGENET1K_V1
model = models.vit_b_16(weights=weights).eval()
preprocess = weights.transforms()

# 推理
img = Image.open("cat.jpg")
batch = preprocess(img).unsqueeze(0)  # [1, 3, 224, 224]

with torch.no_grad():
    logits = model(batch)
    probs = torch.softmax(logits, dim=1)
    top5 = torch.topk(probs, 5)

# 结果
categories = weights.meta["categories"]
for prob, idx in zip(top5.values[0], top5.indices[0]):
    print(f"  {categories[idx]:30s} {prob:.2%}")
# tabby cat                       87.23%
# Egyptian cat                    8.41%
# tiger cat                       3.12%

微调（Transfer Learning）

import torch.nn as nn
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR

# 1. 加载预训练模型
model = models.efficientnet_v2_s(weights=EfficientNet_V2_S_Weights.IMAGENET1K_V1)

# 2. 只替换最后的分类层
num_classes = 10  # 我的数据集的类别数
model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)

# 3. 冻结 backbone（可选）
for param in model.features.parameters():
    param.requires_grad = False  # 固定 backbone

# 4. 只训练分类层
optimizer = AdamW(model.classifier.parameters(), lr=1e-3)
scheduler = CosineAnnealingLR(optimizer, T_max=20)
criterion = nn.CrossEntropyLoss()

# 5. 训练循环
model.train()
for epoch in range(20):
    for images, labels in train_loader:
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    scheduler.step()

# 6. 解冻（微调第2阶段）
for param in model.parameters():
    param.requires_grad = True
optimizer = AdamW(model.parameters(), lr=1e-5)  # 用小 LR！
# 再训练 10 个 epoch...

Part 3: Object Detection

Faster R-CNN

from torchvision.models.detection import (
    fasterrcnn_resnet50_fpn_v2,
    FasterRCNN_ResNet50_FPN_V2_Weights
)

# 预训练模型（COCO 91 个类别）
weights = FasterRCNN_ResNet50_FPN_V2_Weights.COCO_V1
model = fasterrcnn_resnet50_fpn_v2(weights=weights).eval()
preprocess = weights.transforms()

# 推理
img = Image.open("street.jpg")
batch = [preprocess(img)]

with torch.no_grad():
    predictions = model(batch)[0]

# 解析结果
for box, label, score in zip(
    predictions['boxes'], predictions['labels'], predictions['scores']
):
    if score > 0.7:
        category = weights.meta["categories"][label]
        x1, y1, x2, y2 = box.tolist()
        print(f"  {category}: {score:.2%} at ({x1:.0f},{y1:.0f},{x2:.0f},{y2:.0f})")

用自定义数据集训练 Detection

from torchvision.models.detection import fasterrcnn_resnet50_fpn_v2
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor

# 1. 加载预训练模型
model = fasterrcnn_resnet50_fpn_v2(weights="COCO_V1")

# 2. 替换分类头
num_classes = 5 + 1  # 5 个类别 + 背景
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

# 3. 自定义数据集
class MyDetectionDataset(torch.utils.data.Dataset):
    def __getitem__(self, idx):
        img = ...  # 加载图像
        target = {
            "boxes": torch.tensor([[x1,y1,x2,y2], ...], dtype=torch.float32),
            "labels": torch.tensor([1, 3, ...], dtype=torch.int64),
        }
        return img, target

# 4. 训练
model.train()
for images, targets in train_loader:
    loss_dict = model(images, targets)
    # loss_dict: {'loss_classifier', 'loss_box_reg', 'loss_objectness', 'loss_rpn_box_reg'}
    total_loss = sum(loss_dict.values())
    total_loss.backward()
    optimizer.step()
    optimizer.zero_grad()

Part 4: Semantic Segmentation

from torchvision.models.segmentation import (
    deeplabv3_resnet101, DeepLabV3_ResNet101_Weights
)

# DeepLab V3（COCO 21 个类别）
weights = DeepLabV3_ResNet101_Weights.COCO_WITH_VOC_LABELS_V1
model = deeplabv3_resnet101(weights=weights).eval()
preprocess = weights.transforms()

img = Image.open("city.jpg")
batch = preprocess(img).unsqueeze(0)

with torch.no_grad():
    output = model(batch)["out"]  # [1, 21, H, W]
    pred_mask = output.argmax(dim=1)  # [1, H, W] — 每个像素对应的类别

# 可视化
import matplotlib.pyplot as plt
plt.imshow(pred_mask[0].cpu(), cmap="tab20")
plt.title("Semantic Segmentation")
plt.savefig("segmentation.png")

Part 5: 数据集

from torchvision.datasets import (
    CIFAR10, CIFAR100, ImageNet, MNIST,
    FashionMNIST, STL10, Food101, Flowers102,
    CocoDetection, VOCDetection
)

# CIFAR-10（10 个类别, 32x32）
train_set = CIFAR10(root="./data", train=True, download=True, transform=train_transform)

# ImageFolder（自定义数据集）
from torchvision.datasets import ImageFolder

# 目录结构:
# data/train/
#   ├── cat/
#   │   ├── cat001.jpg
#   │   └── cat002.jpg
#   └── dog/
#       ├── dog001.jpg
#       └── dog002.jpg

train_set = ImageFolder(root="data/train", transform=train_transform)
print(train_set.classes)      # ['cat', 'dog']
print(train_set.class_to_idx) # {'cat': 0, 'dog': 1}

train_loader = torch.utils.data.DataLoader(
    train_set, batch_size=32, shuffle=True, num_workers=4, pin_memory=True
)

Part 6: 工具

可视化

from torchvision.utils import make_grid, draw_bounding_boxes, draw_segmentation_masks
import torchvision.transforms.functional as F

# 批量图像网格
grid = make_grid(batch_images, nrow=8, padding=2, normalize=True)
plt.imshow(grid.permute(1, 2, 0))

# Bounding Box 可视化
from torchvision.utils import draw_bounding_boxes
img_with_boxes = draw_bounding_boxes(
    img_tensor,         # uint8, [3, H, W]
    boxes,              # [N, 4]
    labels=["cat", "dog"],
    colors=["red", "blue"],
    width=3,
    font_size=20
)

# Segmentation Mask 可视化
img_with_mask = draw_segmentation_masks(
    img_tensor,
    masks=pred_mask.bool(),
    alpha=0.5,
    colors=["red", "green", "blue"]
)

Feature Extraction（中间层输出）

from torchvision.models.feature_extraction import create_feature_extractor

model = models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)

# 只提取指定层的输出
feature_extractor = create_feature_extractor(model, {
    'layer2': 'mid_features',    # 512 维
    'layer4': 'high_features',   # 2048 维
    'avgpool': 'embedding',       # 2048 维（全局）
})

with torch.no_grad():
    features = feature_extractor(batch)

print(features['mid_features'].shape)   # [1, 512, 28, 28]
print(features['high_features'].shape)  # [1, 2048, 7, 7]
print(features['embedding'].shape)      # [1, 2048, 1, 1]
📝 测验 — torchvision（点击查看！）

Q1. Transforms v2 比 v1 更好的核心原因是什么？ ||可以同时变换图像、Bounding Box 和 Segmentation Mask。v1 中只变换图像，会导致 BBox 错位的问题||

Q2. ImageNet 归一化的 mean 和 std 值是多少？ ||mean: [0.485, 0.456, 0.406], std: [0.229, 0.224, 0.225]。ImageNet 训练数据按 RGB 通道计算出的统计值||

Q3. Transfer Learning 中的 2 阶段微调是什么？ ||第1阶段：冻结 backbone + 只训练分类头（大 LR）。第2阶段：解冻全部层 + 用小 LR 对整体做微调。在保留预训练权重的同时逐步适应新任务||

Q4. Faster R-CNN 的 4 种 loss 是什么？ ||loss_classifier（分类）、loss_box_reg（框回归）、loss_objectness（是否存在物体）、loss_rpn_box_reg（RPN 框回归）||

Q5. create_feature_extractor 的用途是什么？ ||提取模型中间层的输出。无需完整 forward，就能只获取特定层的 feature map，可用于嵌入提取、feature 可视化等场景||

Q6. RandAugment 的 2 个核心参数是什么？ ||num_ops：要应用的 augmentation 操作数量。magnitude：变换强度（0〜30）。与 AutoAugment 不同，仅靠这 2 个参数、无需搜索就能实现强力的 augmentation||

📖 相关系列与推荐阅读

GitHub

测验

Q1：《torchvision 完全指南 — 从图像分类到 Object Detection、Segmentation》一文的主要内容是什么？

从 torchvision 的 transforms v2、预训练模型（ResNet〜ViT）、数据集，到 Object Detection（Faster R-CNN、YOLO）、Segmentation，再到实战微调。用 PyTorch 全面掌握计算机视觉实务。

Q2：Part 1: Transforms v2 — 图像预处理的革新讲了什么？ 基本变换 v2 的核心 — Bounding Box + Mask 同时变换 常用的 Augmentation 配方

Q3：请说明 Part 2: 预训练模型（Model Zoo）的核心概念。 图像分类模型 推理（Inference） 微调（Transfer Learning）

Q4：Part 3: Object Detection 的关键要点有哪些？ Faster R-CNN 用自定义数据集训练 Detection

Q5：Part 6: 工具部分是如何运作的？ 可视化 Feature Extraction（中间层输出） Q1. Transforms v2 比 v1 更好的核心原因是什么？ Q2. ImageNet 归一化的 mean 和 std 值是多少？ Q3. Transfer Learning 中的两阶段微调是什么？ Q4. Faster R-CNN 的 4 种 loss 是什么？ Q5. create_feature_extractor 的用途是什么？

현재 단락 (1/228)

`torchvision` 是 PyTorch 官方的计算机视觉库。图像变换、预训练模型、数据集，直到 Object Detection — CV 所需的几乎一切都包含在内。

작성 글자: 0자원문 글자: 9,464자작성 단락: 0/228