Module de vision par ordinateur PyTorch torchvision

torchvision est une bibliothèque d'extension de l'écosystème PyTorch spécialement conçue pour les tâches de vision par ordinateur. Elle fournit les fonctionnalités principales suivantes :

  1. Modèles pré-entraînés: comprend des implémentations d'architectures CNN classiques (telles que ResNet, VGG, AlexNet, etc.)
  2. Outils de jeux de données: jeux de données visuels courants intégrés (tels que CIFAR10, MNIST, ImageNet, etc.)
  3. Transformations d'images: fournit diverses méthodes de prétraitement d'images et d'augmentation de données
  4. Outils utilitaires: comprend des fonctions auxiliaires telles que le traitement vidéo, les opérations sur les images, etc.
# 安装 torchvision（通常与 PyTorch 一起安装）
pip install torch torchvision

Analyse des composants principaux

1. torchvision.models

Fournit des modèles de vision par ordinateur pré-entraînés pouvant être directement utilisés pour l'apprentissage par transfert :

Exemple

import torchvision.models as models

# Charger le modèle pré-entraîné
resnet18 = models.resnet18(pretrained=True)
alexnet = models.alexnet(pretrained=True)
vgg16 = models.vgg16(pretrained=True)

Liste des modèles courants :

Nom du modèle Scénario applicable Nombre de paramètres Précision Top-1
ResNet Classification générale d'images 11M-60M 69%-80%
VGG Extraction de caractéristiques 138M 71.3%
MobileNet Application mobile 3.4M 70.6%
EfficientNet Modèle efficace 5M-66M 77%-84%

2. torchvision.datasets

Jeux de données de vision par ordinateur courants intégrés pour simplifier le processus de chargement des données :

Exemple

from torchvision import datasets

# Charger le jeu de données CIFAR10
train_data = datasets.CIFAR10(
    root='data',
    train=True,
    download=True,
    transform=transforms.ToTensor()
)

# Charger le jeu de données MNIST
test_data = datasets.MNIST(
    root='data',
    train=False,
    download=True
)

Types de jeux de données pris en charge :

Exemple

graph TD
A[torchvision.datasets] --> B[Jeux de données de classification]
A --> C[Jeux de données de détection]
A --> D[Jeux de données de segmentation]
    B --> B1[CIFAR10/100]
    B --> B2[MNIST/FashionMNIST]
    B --> B3[ImageNet]
    C --> C1[COCO]
    C --> C2[VOC]
    D --> D1[Cityscapes]

3. torchvision.transforms

Outils essentiels pour le prétraitement d'images et l'augmentation de données :

Exemple

from torchvision import transforms

# Définir le pipeline de transformation d'images
transform = transforms.Compose([
    transforms.Resize(256),          # Redimensionner
    transforms.CenterCrop(224),       # Recadrage central
    transforms.ToTensor(),           # Convertir en tenseur
    transforms.Normalize(             # Normaliser
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

Classification des méthodes de transformation courantes :

Catégorie Exemples de méthodes Objectif
Transformations géométriques RandomRotation, RandomResizedCrop Augmenter l'invariance de position
Transformations de couleur ColorJitter, Grayscale Améliorer la robustesse aux couleurs
Flou/Bruit GaussianBlur, RandomErasing Éviter le surapprentissage
Transformations combinées RandomApply, RandomChoice Stratégie de combinaison flexible

Exemple pratique : pipeline de classification d'images

1. Préparation des données

Exemple

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

# Définir les transformations de données
train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(10),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# Charger le jeu de données
train_set = datasets.CIFAR10(
    root='./data',
    train=True,
    download=True,
    transform=train_transform
)

# Créer le chargeur de données
train_loader = DataLoader(
    train_set,
    batch_size=32,
    shuffle=True
)

2. Entraînement du modèle

Exemple

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

# Utiliser un modèle pré-entraîné
model = models.resnet18(pretrained=True)

# Modifier la dernière couche (adapter aux 10 classes de CIFAR10)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 10)

# Définir la fonction de perte et l'optimiseur
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# Boucle d'entraînement
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()

Fonctionnalités avancées

1. Jeu de données personnalisé

Exemple

from torchvision.datasets import VisionDataset

class CustomDataset(VisionDataset):
    def __init__(self, root, transform=None):
        super().__init__(root, transform=transform)
        # Implémenter __getitem__ et __len__
       
    def __getitem__(self, index):
        # Renvoyer le tuple (image, target)
        pass
       
    def __len__(self):
        # Renvoyer la taille du jeu de données
        pass

2. Export et déploiement du modèle

Exemple

# Exporter au format ONNX
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["input"],
    output_names=["output"]
)

Recommandations de bonnes pratiques

1. Stratégie d'augmentation des données：

  • Utilisez des transformations aléatoires pour augmenter les données pendant l'entraînement
  • Utilisez uniquement des transformations déterministes lors de la validation/du test

2. Conseils pour l'apprentissage par transfert：

Exemple

# Figer tous les paramètres sauf la dernière couche
for param in model.parameters():
    param.requires_grad = False
model.fc.requires_grad = True

3. Optimisation des performances：

  • Utilisernum_workersparamètre pour accélérer le chargement des données
  • Pour les grands jeux de données, envisagez d'utiliserDatasetsous-ensemble de

4. Erreurs courantes：

  • Oublier d'appelerzero_grad()
  • Confusiontrain()eteval()mode
  • La forme du tenseur d'image ne correspond pas aux exigences du modèle (doit être C×H×W)
Autres extensions