Exemple PyTorch - Classification d'images

La classification d'images est l'une des tâches les plus fondamentales en vision par ordinateur. Son objectif est de permettre aux ordinateurs de reconnaître le contenu principal d'une image et de la classer dans des classes prédéfinies. Par exemple, identifier si une image contient un chat ou un chien.

Application de l'apprentissage profond à la classification d'images

Les modèles d'apprentissage profond, les réseaux de neurones convolutifs (CNN), sont devenus la solution dominante pour les tâches de classification d'images. En tant que framework d'apprentissage profond, PyTorch fournit une chaîne d'outils complète pour construire et entraîner des modèles CNN.

Processus du projet

Un projet typique de classification d'images comprend les étapes suivantes :

Préparation et prétraitement des données
Modèle
Entraînement du modèle
Évaluation du modèle
Application du modèle

Préparation de l'environnement et chargement des données

Installer les bibliothèques requises

# 安装PyTorch和torchvision
!pip install torch torchvision

Charger les jeux de données courants

La torchvision de PyTorch fournit plusieurs jeux de données courants, tels que CIFAR-10, MNIST, etc.

Exemple

import torch
import torchvision
import torchvision.transforms as transforms

# Définir les transformations de données
transform = transforms.Compose([
transforms.ToTensor(), # Convertir une image PIL en Tensor
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) # Normalisation
])

# Charger le jeu d'entraînement CIFAR-10
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
shuffle=True, num_workers=2)

# Charger le jeu de test CIFAR-10
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
shuffle=False, num_workers=2)

# Définir les noms des classes
classes = ('plane', 'car', 'bird', 'cat', 'deer',
'dog', 'frog', 'horse', 'ship', 'truck')

Construire un modèle de réseau de neurones convolutif

Structure de base d'un CNN

Un CNN typique se compose de couches convolutives, de couches de pooling et de couches entièrement connectées.

Exemple

import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
# Couche de convolution 1 : entrée 3 canaux (RVB), sortie 6 canaux, noyau 5x5
self.conv1 = nn.Conv2d(3, 6, 5)
# Couche de pooling : fenêtre 2x2, pas 2
self.pool = nn.MaxPool2d(2, 2)
# Couche de convolution 2 : entrée 6 canaux, sortie 16 canaux, noyau 5x5
self.conv2 = nn.Conv2d(6, 16, 5)
# Couche entièrement connectée 1 : entrée 16*5*5, sortie 120
self.fc1 = nn.Linear(16 * 5 * 5, 120)
# Couche entièrement connectée 2 : entrée 120, sortie 84
self.fc2 = nn.Linear(120, 84)
# Couche entièrement connectée 3 : entrée 84, sortie 10 (correspondant à 10 classes)
self.fc3 = nn.Linear(84, 10)

def forward(self, x):
# Première couche de convolution + ReLU + pooling
x = self.pool(F.relu(self.conv1(x)))
# Deuxième couche de convolution + ReLU + pooling
x = self.pool(F.relu(self.conv2(x)))
# Aplatir les cartes de caractéristiques
x = x.vue(-1, 16 * 5 * 5)
# Couche entièrement connectée + ReLU
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
# Couche de sortie
x = self.fc3(x)
return x

# Instancier le réseau
net = Net()

Visualisation de la structure du modèle

Entraîner le modèle

Définir la fonction de perte et l'optimiseur

Exemple

import torch.optim as optim

# Fonction de perte d'entropie croisée
criterion = nn.CrossEntropyLoss()
# Optimiseur de descente de gradient stochastique
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

Boucle d'entraînement

Exemple

for epoch in range(10): # Entraîner pendant 10 époques
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
# Obtenir les données d'entrée
inputs, labels = data

# Remettre les gradients à zéro
optimizer.zero_grad()

# Propagation avant
outputs = net(inputs)
# Calculer la perte
loss = criterion(outputs, labels)
# Rétropropagation
loss.backward()
# Mettre à jour les poids
optimizer.step()

# Afficher les statistiques
running_loss += loss.item()
if i % 2000 == 1999: # Afficher tous les 2000 mini-lots
print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
running_loss = 0.0

print('Entraînement terminé')

Évaluation du modèle

Calcul de la précision sur l'ensemble de test

Exemple

correct = 0
total = 0
with torch.no_grad(): # Ne pas calculer les gradients
for data in testloader:
images, labels = data
outputs = net(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()

print(f'Précision sur les images de test : {100 * correct / total:.2f}%')

Analyse de la précision par classe

Exemple

class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))
with torch.no_grad():
for data in testloader:
images, labels = data
outputs = net(images)
_, predicted = torch.max(outputs, 1)
c = (predicted == labels).squeeze()
for i in range(4):
label = labels[i]
class_correct[label] += c[i].item()
class_total[label] += 1

for i in range(10):
print(f'Précision de {classes[i]:5s}: {100 * class_correct[i] / class_total[i]:.2f}%')

Sauvegarde et chargement du modèle

Enregistrer le modèle entraîné

Exemple

# Enregistrer les paramètres du modèle
PATH = './cifar_net.pth'
torch.save(net.state_dict(), PATH)

Charger le modèle pour la prédiction

Exemple

# Charger le modèle
net = Net()
net.load_state_dict(torch.load(PATH))

# Utilisez le modèle pour faire des prédictions
outputs = net(images)
_, predicted = torch.max(outputs, 1)
print('Prédit : ', ' '.join(f'{classes[predicted[j]]:5s}' for j in range(4)))

Autres extensions