PyTorch-Entwicklungsmuster
Idiomatische PyTorch-Muster und Best Practices für den Aufbau robuster, effizienter und reproduzierbarer Deep-Learning-Anwendungen.
Wann aktivieren
- Schreiben neuer PyTorch-Modelle oder Trainingsskripte
- Überprüfen von Deep-Learning-Code
- Debuggen von Trainingsschleifen oder Datenpipelines
- Optimieren der GPU-Speichernutzung oder Trainingsgeschwindigkeit
- Einrichten reproduzierbarer Experimente
Kernprinzipien
1. Geräteagnostischer Code
Schreiben Sie Code, der ohne hartcodierte Geräte sowohl auf CPU als auch auf GPU funktioniert.
# Gut: Geräteagnostisch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MyModel().to(device)
data = data.to(device)
# Schlecht: Hartcodiertes Gerät
model = MyModel().cuda() # Absturz, wenn keine GPU
data = data.cuda()
2. Reproduzierbarkeit zuerst
Setzen Sie alle Zufallsinitialisierungen für reproduzierbare Ergebnisse.
# Gut: Vollständige Reproduzierbarkeitseinrichtung
def set_seed(seed: int = 42) -> None:
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
np.random.seed(seed)
random.seed(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# Schlecht: Keine Seed-Kontrolle
model = MyModel() # Andere Gewichte bei jedem Durchlauf
3. Explizite Shape-Verwaltung
Dokumentieren und überprüfen Sie stets Tensor-Shapes.
# Gut: Mit Shape-Kommentaren versehener Forward-Durchlauf
def forward(self, x: torch.Tensor) -> torch.Tensor:
# x: (batch_size, channels, height, width)
x = self.conv1(x) # -> (batch_size, 32, H, W)
x = self.pool(x) # -> (batch_size, 32, H//2, W//2)
x = x.view(x.size(0), -1) # -> (batch_size, 32*H//2*W//2)
return self.fc(x) # -> (batch_size, num_classes)
# Schlecht: Keine Shape-Verfolgung
def forward(self, x):
x = self.conv1(x)
x = self.pool(x)
x = x.view(x.size(0), -1) # Welche Größe ist das?
return self.fc(x) # Funktioniert das überhaupt?
Modellarchitektur-Muster
Saubere nn.Module-Struktur
# Gut: Gut organisierte Module
class ImageClassifier(nn.Module):
def __init__(self, num_classes: int, dropout: float = 0.5) -> None:
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=3, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
nn.MaxPool2d(2),
)
self.classifier = nn.Sequential(
nn.Dropout(dropout),
nn.Linear(64 * 16 * 16, num_classes),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.features(x)
x = x.view(x.size(0), -1)
return self.classifier(x)
# Schlecht: Alles in forward
class ImageClassifier(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x):
x = F.conv2d(x, weight=self.make_weight()) # Erzeugt bei jedem Aufruf Gewichte!
return x
Ordnungsgemäße Gewichtsinitialisierung
# Gut: Explizite Initialisierung
def _init_weights(self, module: nn.Module) -> None:
if isinstance(module, nn.Linear):
nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
if module.bias is not None:
nn.init.zeros_(module.bias)
elif isinstance(module, nn.Conv2d):
nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
elif isinstance(module, nn.BatchNorm2d):
nn.init.ones_(module.weight)
nn.init.zeros_(module.bias)
model = MyModel()
model.apply(model._init_weights)
Trainingsschleifen-Muster
Standard-Trainingsschleife
# Gut: Vollständige Trainingsschleife mit Best Practices
def train_one_epoch(
model: nn.Module,
dataloader: DataLoader,
optimizer: torch.optim.Optimizer,
criterion: nn.Module,
device: torch.device,
scaler: torch.amp.GradScaler | None = None,
) -> float:
model.train() # Immer Trainingsmodus setzen
total_loss = 0.0
for batch_idx, (data, target) in enumerate(dataloader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad(set_to_none=True) # Effizienter als zero_grad()
# Mixed-Precision-Training
with torch.amp.autocast("cuda", enabled=scaler is not None):
output = model(data)
loss = criterion(output, target)
if scaler is not None:
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
total_loss += loss.item()
return total_loss / len(dataloader)
Validierungsschleife
# Gut: Korrekte Evaluierung
@torch.no_grad() # Effizienter als Umwickeln mit torch.no_grad()-Block
def evaluate(
model: nn.Module,
dataloader: DataLoader,
criterion: nn.Module,
device: torch.device,
) -> tuple[float, float]:
model.eval() # Immer Evaluierungsmodus setzen – deaktiviert Dropout, verwendet laufende BN-Statistiken
total_loss = 0.0
correct = 0
total = 0
for data, target in dataloader:
data, target = data.to(device), target.to(device)
output = model(data)
total_loss += criterion(output, target).item()
correct += (output.argmax(1) == target).sum().item()
total += target.size(0)
return total_loss / len(dataloader), correct / total
Datenpipeline-Muster
Benutzerdefinierter Dataset
# Gut: Sauberer Dataset mit Typannotationen
class ImageDataset(Dataset):
def __init__(
self,
image_dir: str,
labels: dict[str, int],
transform: transforms.Compose | None = None,
) -> None:
self.image_paths = list(Path(image_dir).glob("*.jpg"))
self.labels = labels
self.transform = transform
def __len__(self) -> int:
return len(self.image_paths)
def __getitem__(self, idx: int) -> tuple[torch.Tensor, int]:
img = Image.open(self.image_paths[idx]).convert("RGB")
label = self.labels[self.image_paths[idx].stem]
if self.transform:
img = self.transform(img)
return img, label
Effiziente DataLoader-Konfiguration
# Gut: Optimierter DataLoader
dataloader = DataLoader(
dataset,
batch_size=32,
shuffle=True, # Mischen für Training
num_workers=4, # Paralleles Datenladen
pin_memory=True, # Schnellerer CPU->GPU-Transfer
persistent_workers=True, # Arbeiter zwischen Epochen am Leben halten
drop_last=True, # Konsistente Batch-Größen für BatchNorm
)
# Schlecht: Langsame Standardwerte
dataloader = DataLoader(dataset, batch_size=32) # num_workers=0, kein pin_memory
Benutzerdefiniertes Collate für Daten variabler Länge
# Gut: Sequenzen in collate_fn auffüllen
def collate_fn(batch: list[tuple[torch.Tensor, int]]) -> tuple[torch.Tensor, torch.Tensor]:
sequences, labels = zip(*batch)
# Auf maximale Länge im Batch auffüllen
padded = nn.utils.rnn.pad_sequence(sequences, batch_first=True, padding_value=0)
return padded, torch.tensor(labels)
dataloader = DataLoader(dataset, batch_size=32, collate_fn=collate_fn)
Checkpointing-Muster
Checkpoints speichern und laden
# Gut: Vollständiger Checkpoint mit gesamten Trainingsstatus
def save_checkpoint(
model: nn.Module,
optimizer: torch.optim.Optimizer,
epoch: int,
loss: float,
path: str,
) -> None:
torch.save({
"epoch": epoch,
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"loss": loss,
}, path)
def load_checkpoint(
path: str,
model: nn.Module,
optimizer: torch.optim.Optimizer | None = None,
) -> dict:
checkpoint = torch.load(path, map_location="cpu", weights_only=True)
model.load_state_dict(checkpoint["model_state_dict"])
if optimizer:
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
return checkpoint
# Schlecht: Nur Modellgewichte speichern (Training nicht fortsetzbar)
torch.save(model.state_dict(), "model.pt")
Leistungsoptimierung
Mixed-Precision-Training
# Gut: AMP mit GradScaler
scaler = torch.amp.GradScaler("cuda")
for data, target in dataloader:
with torch.amp.autocast("cuda"):
output = model(data)
loss = criterion(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)
Gradienten-Checkpointing für große Modelle
# Gut: Rechenleistung gegen Speicher tauschen
from torch.utils.checkpoint import checkpoint
class LargeModel(nn.Module):
def forward(self, x: torch.Tensor) -> torch.Tensor:
# Aktivierungen während des Rückwärtsdurchlaufs neu berechnen, um Speicher zu sparen
x = checkpoint(self.block1, x, use_reentrant=False)
x = checkpoint(self.block2, x, use_reentrant=False)
return self.head(x)
torch.compile für Geschwindigkeit
# Gut: Modell für schnellere Ausführung kompilieren (PyTorch 2.0+)
model = MyModel().to(device)
model = torch.compile(model, mode="reduce-overhead")
# Modi: "default" (sicher), "reduce-overhead" (schneller), "max-autotune" (am schnellsten)
Kurzreferenz: PyTorch-Idiome
| Idiom | Beschreibung |
|---|---|
model.train() / model.eval() | Modus vor Training/Evaluierung immer setzen |
torch.no_grad() | Gradienten für Inferenz deaktivieren |
optimizer.zero_grad(set_to_none=True) | Effizientere Gradientenlöschung |
.to(device) | Geräteagnostische Tensor-/Modellplatzierung |
torch.amp.autocast | Mixed Precision für 2x Geschwindigkeit |
pin_memory=True | Schnellerer CPU→GPU-Datentransfer |
torch.compile | JIT-Kompilierung für Geschwindigkeit (2.0+) |
weights_only=True | Sicheres Modellladen |
torch.manual_seed | Reproduzierbare Experimente |
gradient_checkpointing | Rechenleistung gegen Speicher tauschen |
Zu vermeidende Anti-Patterns
# Schlecht: Vergessen von model.eval() während der Validierung
model.train()
with torch.no_grad():
output = model(val_data) # Dropout immer noch aktiv! BatchNorm verwendet Batch-Statistiken!
# Gut: Immer Evaluierungsmodus setzen
model.eval()
with torch.no_grad():
output = model(val_data)
# Schlecht: In-Place-Operationen, die Autograd unterbrechen
x = F.relu(x, inplace=True) # Kann Gradientenberechnung unterbrechen
x += residual # In-Place-Add unterbricht Autograd-Graph
# Gut: Out-of-Place-Operationen
x = F.relu(x)
x = x + residual
# Schlecht: Wiederholtes Verschieben von Daten auf die GPU innerhalb der Trainingsschleife
for data, target in dataloader:
model = model.cuda() # Verschiebt Modell bei JEDER Iteration!
# Gut: Modell einmal vor der Schleife verschieben
model = model.to(device)
for data, target in dataloader:
data, target = data.to(device), target.to(device)
# Schlecht: Verwenden von .item() vor backward
loss = criterion(output, target).item() # Trennt vom Graphen!
loss.backward() # Fehler: kein Backprop durch .item() möglich
# Gut: .item() nur für Protokollierung aufrufen
loss = criterion(output, target)
loss.backward()
print(f"Loss: {loss.item():.4f}") # .item() nach backward ist in Ordnung
# Schlecht: torch.save nicht korrekt verwenden
torch.save(model, "model.pt") # Speichert gesamtes Modell (zerbrechlich, nicht portabel)
# Gut: state_dict speichern
torch.save(model.state_dict(), "model.pt")
Denken Sie daran: PyTorch-Code sollte geräteagnostisch, reproduzierbar und speicherbewusst sein. Im Zweifelsfall mit torch.profiler profilieren und GPU-Speicher mit torch.cuda.memory_summary() überprüfen.