Fonction PyTorch torch.nn.Module
Manuel de référence torch.nn de PyTorch
torch.nn.ModuleC'est la classe de base pour tous les modules de réseaux de neurones dans PyTorch.
Tous les modèles de réseau personnalisés doivent hériter de cette classe, car elle fournit la gestion des paramètres, la migration vers un appareil, l'enregistrement du modèle et d'autres fonctions.
Définition de classe
torch.nn.Module
Attributs principaux
parameters(): Renvoie tous les paramètres apprenables du modèle.named_parameters(): Renvoie un itérateur de noms et de valeurs des paramètres.children(): Renvoie les sous-modules du modèle.named_children(): Renvoie un itérateur de noms de sous-modules et de modules.modules(): Renvoie tous les modules.state_dict(): Renvoie un dictionnaire contenant tous les paramètres.
Exemples d'utilisation
Exemple 1: Création d'un module personnalisé
Héritez de nn.Module pour créer un réseau personnalisé :
Exemple
import torch
import torch.nn as nn
class SimpleNet(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super(SimpleNet, self).__init__()
# Définir les couches du réseau
self.fc1 = nn.Linear(input_dim, hidden_dim)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(hidden_dim, output_dim)
def forward(self, x):
# Définir la propagation avant
x = self.fc1(x)
x = self.relu(x)
x = self.fc2(x)
return x
# Instancier le modèle
model = SimpleNet(input_dim=784, hidden_dim=256, output_dim=10)
# Test
x = torch.randn(32, 784)
output = model(x)
print(Structure du modèle :)
print(model)
print(nForme d'entrée :, x.shape)
print(Forme de sortie :, output.shape)
import torch.nn as nn
class SimpleNet(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super(SimpleNet, self).__init__()
# Définir les couches du réseau
self.fc1 = nn.Linear(input_dim, hidden_dim)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(hidden_dim, output_dim)
def forward(self, x):
# Définir la propagation avant
x = self.fc1(x)
x = self.relu(x)
x = self.fc2(x)
return x
# Instancier le modèle
model = SimpleNet(input_dim=784, hidden_dim=256, output_dim=10)
# Test
x = torch.randn(32, 784)
output = model(x)
print(Structure du modèle :)
print(model)
print(nForme d'entrée :, x.shape)
print(Forme de sortie :, output.shape)
Exemple 2 : Gestion des paramètres
Accès et gestion des paramètres du modèle :
Exemple
import torch
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 16, 3)
self.bn1 = nn.BatchNorm2d(16)
self.fc = nn.Linear(16 * 14 * 14, 10)
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
model = Net()
# Compter le nombre de paramètres
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(Nombre total de paramètres :, total_params)
print(Paramètres entraînables :, trainable_params)
# Accéder à des paramètres spécifiques
print(nForme du poids de conv1 :, model.conv1.weight.shape)
print(Forme du biais de fc :, model.fc.bias.shape)
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 16, 3)
self.bn1 = nn.BatchNorm2d(16)
self.fc = nn.Linear(16 * 14 * 14, 10)
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
model = Net()
# Compter le nombre de paramètres
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(Nombre total de paramètres :, total_params)
print(Paramètres entraînables :, trainable_params)
# Accéder à des paramètres spécifiques
print(nForme du poids de conv1 :, model.conv1.weight.shape)
print(Forme du biais de fc :, model.fc.bias.shape)
Exemple 3: Sauvegarde et chargement du modèle
Sauvegarder et charger le modèle:
Exemple
import torch
import torch.nn as nn
import tempfile
import os
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc = nn.Linear(10, 5)
def forward(self, x):
return self.fc(x)
model = Net()
# Enregistrer l'intégralité du modèle
with tempfile.NamedTemporaryFile(delete=False, suffix='.pth') as f:
torch.save(model, f.name)
path_full = f.name
# Sauvegarder state_dict (méthode recommandée)
with tempfile.NamedTemporaryFile(delete=False, suffix='.pth') as f:
torch.save(model.state_dict(), f.name)
path_state = f.name
# Charger le modèle
loaded_model = Net()
loaded_model.load_state_dict(torch.load(path_state))
loaded_model.eval()
# Tester le modèle chargé
x = torch.randn(2, 10)
output1 = model(x)
output2 = loaded_model(x)
print(Sortie d'origine :, output1[0].tolist())
print(Sortie après chargement :, output2[0].tolist())
# Nettoyage
os.remove(path_full)
os.remove(path_state)
import torch.nn as nn
import tempfile
import os
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc = nn.Linear(10, 5)
def forward(self, x):
return self.fc(x)
model = Net()
# Enregistrer l'intégralité du modèle
with tempfile.NamedTemporaryFile(delete=False, suffix='.pth') as f:
torch.save(model, f.name)
path_full = f.name
# Sauvegarder state_dict (méthode recommandée)
with tempfile.NamedTemporaryFile(delete=False, suffix='.pth') as f:
torch.save(model.state_dict(), f.name)
path_state = f.name
# Charger le modèle
loaded_model = Net()
loaded_model.load_state_dict(torch.load(path_state))
loaded_model.eval()
# Tester le modèle chargé
x = torch.randn(2, 10)
output1 = model(x)
output2 = loaded_model(x)
print(Sortie d'origine :, output1[0].tolist())
print(Sortie après chargement :, output2[0].tolist())
# Nettoyage
os.remove(path_full)
os.remove(path_state)
Exemple 4 : Migration de périphérique
Migrer le modèle entre différents périphériques :
Exemple
import torch
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc = nn.Linear(10, 5)
def forward(self, x):
return self.fc(x)
model = Net()
# Vérifier le périphérique actuel
print(Périphérique du paramètre :, model.fc.weight.device)
# Déplacer vers le GPU (si disponible)
if torch.cuda.is_available():
model = model.cuda()
print(Après déplacement vers le GPU :, model.fc.weight.device)
# Revenir au CPU
model = model.cpu()
print(Après retour au CPU :, model.fc.weight.device)
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc = nn.Linear(10, 5)
def forward(self, x):
return self.fc(x)
model = Net()
# Vérifier le périphérique actuel
print(Périphérique du paramètre :, model.fc.weight.device)
# Déplacer vers le GPU (si disponible)
if torch.cuda.is_available():
model = model.cuda()
print(Après déplacement vers le GPU :, model.fc.weight.device)
# Revenir au CPU
model = model.cpu()
print(Après retour au CPU :, model.fc.weight.device)
Exemple 5 : Initialisation avec apply
Utilisez apply pour initialiser récursivement :
Exemple
import torch
import torch.nn as nn
def init_weights(module):
if isinstance(module, nn.Linear):
nn.init.xavier_uniform_(module.weight)
if module.bias is not None:
nn.init.zeros_(module.bias)
elif isinstance(module, nn.Conv2d):
nn.init.kaiming_normal_(module.weight)
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(10, 20)
self.fc2 = nn.Linear(20, 5)
def forward(self, x):
return self.fc2(self.fc1(x))
model = Net()
model.apply(init_weights)
print(fc1 poids:, model.fc1.weight[0, :3].tolist())
print(fc2 poids:, model.fc2.weight[0, :3].tolist())
import torch.nn as nn
def init_weights(module):
if isinstance(module, nn.Linear):
nn.init.xavier_uniform_(module.weight)
if module.bias is not None:
nn.init.zeros_(module.bias)
elif isinstance(module, nn.Conv2d):
nn.init.kaiming_normal_(module.weight)
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(10, 20)
self.fc2 = nn.Linear(20, 5)
def forward(self, x):
return self.fc2(self.fc1(x))
model = Net()
model.apply(init_weights)
print(fc1 poids:, model.fc1.weight[0, :3].tolist())
print(fc2 poids:, model.fc2.weight[0, :3].tolist())
Exemple 6: Structure de réseau complexe
Construisez un réseau avec des branches:
Exemple
import torch
import torch.nn as nn
class BranchNet(nn.Module):
def __init__(self):
super(BranchNet, self).__init__()
# Colonne vertébrale
self.shared = nn.Linear(10, 20)
# Branche
self.branch_a = nn.Linear(20, 5)
self.branch_b = nn.Linear(20, 3)
def forward(self, x):
feat = self.shared(x)
out_a = self.branch_a(feat)
out_b = self.branch_b(feat)
return out_a, out_b
model = BranchNet()
x = torch.randn(4, 10)
out_a, out_b = model(x)
print(Sortie de la branche A:, out_a.shape)
print(Sortie de la branche B:, out_b.shape)
import torch.nn as nn
class BranchNet(nn.Module):
def __init__(self):
super(BranchNet, self).__init__()
# Colonne vertébrale
self.shared = nn.Linear(10, 20)
# Branche
self.branch_a = nn.Linear(20, 5)
self.branch_b = nn.Linear(20, 3)
def forward(self, x):
feat = self.shared(x)
out_a = self.branch_a(feat)
out_b = self.branch_b(feat)
return out_a, out_b
model = BranchNet()
x = torch.randn(4, 10)
out_a, out_b = model(x)
print(Sortie de la branche A:, out_a.shape)
print(Sortie de la branche B:, out_b.shape)
FAQ
Q1: Pourquoi faut-il appeler super().__init__()?
Il appelle la méthode d'initialisation de la classe parente, ce qui garantit que le module enregistre correctement les paramètres.
<h3 Q2: Comment visualiser la structure du modèle?Affichez directement print(model) ou utilisez la fonction summary de torchvision.
<h3 Q3 : Quelle est la différence entre parameters() et modules() ?parameters() renvoie les paramètres apprenables, tandis que modules() renvoie tous les modules.
Cas d'utilisation
nn.ModuleC'est la base pour construire tous les réseaux de neurones personnalisés. Les principaux scénarios d'application incluent :
- Définition de la structure du réseau: Tout réseau de neurones personnalisé
- Gestion du modèle: Sauvegarde, chargement, migration
- Gestion des paramètres: Mises à jour de l'optimiseur, calcul du gradient
Remarque : Tous les Modules personnalisés doivent implémenter la méthode forward.
Autres extensions