PyTorch torchvision Computer-Vision-Modul

torchvision ist eine Erweiterungsbibliothek im PyTorch-Ökosystem speziell für Computer-Vision-Aufgaben. Sie bietet die folgenden Kernfunktionen:

  1. Vortrainierte Modelle: Enthält Implementierungen klassischer CNN-Architekturen (wie ResNet, VGG, AlexNet usw.)
  2. Dataset-Werkzeuge: Integrierte gängige visuelle Datensätze (wie CIFAR10, MNIST, ImageNet usw.)
  3. Bildtransformationen: Bietet verschiedene Methoden zur Bildvorverarbeitung und Daten-Augmentierung
  4. Hilfswerkzeuge: Enthält Hilfsfunktionen wie Videoverarbeitung, Bildoperationen usw.
# 安装 torchvision（通常与 PyTorch 一起安装）
pip install torch torchvision

Analyse der Kernkomponenten

1. torchvision.models

Bietet vortrainierte Computer-Vision-Modelle, die direkt für Transfer Learning verwendet werden können:

Beispiel

import torchvision.models as models

# Vortrainiertes Modell laden
resnet18 = models.resnet18(pretrained=True)
alexnet = models.alexnet(pretrained=True)
vgg16 = models.vgg16(pretrained=True)

Liste gängiger Modelle:

Modellname Anwendungsfall Parameteranzahl Top-1-Genauigkeit
ResNet Allgemeine Bildklassifizierung 11M-60M 69%-80%
VGG Merkmalsextraktion 138M 71.3%
MobileNet Mobile Anwendung 3.4M 70.6%
EfficientNet Effizientes Modell 5M-66M 77%-84%

2. torchvision.datasets

Integrierte gängige Computer-Vision-Datensätze zur Vereinfachung des Datenladevorgangs:

Beispiel

from torchvision import datasets

# CIFAR10-Datensatz laden
train_data = datasets.CIFAR10(
    root='data',
    train=True,
    download=True,
    transform=transforms.ToTensor()
)

# MNIST-Datensatz laden
test_data = datasets.MNIST(
    root='data',
    train=False,
    download=True
)

Unterstützte Datensatztypen:

Beispiel

graph TD
A[torchvision.datasets] --> B[Klassifizierungsdatensätze]
A --> C[Detektionsdatensätze]
A --> D[Segmentierungsdatensätze]
    B --> B1[CIFAR10/100]
    B --> B2[MNIST/FashionMNIST]
    B --> B3[ImageNet]
    C --> C1[COCO]
    C --> C2[VOC]
    D --> D1[Cityscapes]

3. torchvision.transforms

Zentrale Werkzeuge für Bildvorverarbeitung und Datenaugmentierung:

Beispiel

from torchvision import transforms

# Bildtransformations-Pipeline definieren
transform = transforms.Compose([
    transforms.Resize(256),          # Größe ändern
    transforms.CenterCrop(224),       # Zentrierter Zuschnitt
    transforms.ToTensor(),           # In Tensor umwandeln
    transforms.Normalize(             # Normalisieren
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

Klassifizierung gängiger Transformationsmethoden:

Kategorie Beispielmethoden Zweck
Geometrische Transformationen RandomRotation, RandomResizedCrop Positionsinvarianz erhöhen
Farbtransformationen ColorJitter, Grayscale Farbrobustheit verbessern
Unschärfe/Rauschen GaussianBlur, RandomErasing Overfitting verhindern
Kombinierte Transformationen RandomApply, RandomChoice Flexible Kombinationsstrategie

Praktisches Beispiel: Bildklassifizierungs-Pipeline

1. Datenvorbereitung

Beispiel

import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# Datentransformationen definieren
train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(10),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# Datensatz laden
train_set = datasets.CIFAR10(
    root='./data',
    train=True,
    download=True,
    transform=train_transform
)

# Data Loader erstellen
train_loader = DataLoader(
    train_set,
    batch_size=32,
    shuffle=True
)

2. Modelltraining

Beispiel

import torch.nn as nn
import torch.optim as optim

# Vortrainiertes Modell verwenden
model = models.resnet18(pretrained=True)

# Letzte Schicht ändern (an die 10 Klassen von CIFAR10 anpassen)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 10)

# Verlustfunktion und Optimierer definieren
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# Trainingsschleife
for epoch in range(10):
    for images, labels in train_loader:
        outputs = model(images)
        loss = criterion(outputs, labels)
       
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

Erweiterte Funktionen

1. Benutzerdefinierter Datensatz

Beispiel

from torchvision.datasets import VisionDataset

class CustomDataset(VisionDataset):
    def __init__(self, root, transform=None):
        super().__init__(root, transform=transform)
        # Implementiere __getitem__ und __len__
       
    def __getitem__(self, index):
        # Gib (image, target)-Tupel zurück
        pass
       
    def __len__(self):
        # Datensatzgröße zurückgeben
        pass

2. Modell-Export und Bereitstellung

Beispiel

# In das ONNX-Format exportieren
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["input"],
    output_names=["output"]
)

Best-Practice-Empfehlungen

1. Strategie zur Datenaugmentierung：

  • Verwenden Sie zufällige Transformationen, um Daten während des Trainings zu augmentieren
  • Verwenden Sie nur deterministische Transformationen während der Validierung/Testphase

2. Tipps zum Transferlernen：

Beispiel

# Alle Parameter außer der letzten Schicht einfrieren
for param in model.parameters():
    param.requires_grad = False
model.fc.requires_grad = True

3. Leistungsoptimierung：

  • Verwenden Sienum_workersParameter zur Beschleunigung des Datenladens
  • Bei großen Datensätzen sollten Sie die Verwendung von ... in Betracht ziehenDatasetTeilmenge von

4. Häufige Fehler：

  • Vergessen aufzurufenzero_grad()
  • Verwirrttrain()undeval()mode
  • Die Form des Bildtensors entspricht nicht den Modellanforderungen (sollte C×H×W sein)
Andere Erweiterungen