PyTorch TorchScript / ONNX-Export
Nach Abschluss des Modelltrainings muss das Modell in die Produktionsumgebung bereitgestellt werden. PyTorch bietet mehrere Modell-Exportmethoden, wobeiTorchScriptundONNXdie beiden am häufigsten verwendeten Formate sind.
Dieser Abschnitt erläutert die Prinzipien, Verwendung und bewährten Verfahren dieser beiden Exportmethoden.
Modell-Export ist der Prozess der Konvertierung eines PyTorch-Modells in ein Format, das auf verschiedenen Plattformen und Frameworks ausgeführt werden kann. Dies ist entscheidend für Szenarien wie Modellbereitstellung, mobile Inferenz und frameworkübergreifende Migration.
1. TorchScript-Grundlagen
1.1 Was ist TorchScript
TorchScript ist das Serialisierungsformat von PyTorch, das Python-Code in eigenständigen C++-Code für virtuelle Maschinen umwandelt. TorchScript-Programme können in Umgebungen ohne Python-Interpreter ausgeführt werden.
Hauptmerkmale von TorchScript:
- Wandelt dynamische Graphen in statische Graphen um
- Unterstützt eine Teilmenge der Python-Syntax
- Kann in einer C++-Umgebung ausgeführt werden
- Bewahrt die Struktur und die Parameter des Modells
1.2 Zwei Konvertierungsmethoden
TorchScript bietet zwei Möglichkeiten, Modelle in TorchScript zu konvertieren:
- TorchScript-Tracing (Tracing): Zeichnet Operationen durch Ausführen des Modells auf und erzeugt einen statischen Berechnungsgraphen
- TorchScript-Scripting (Scripting): Analysiert Python-Code direkt und kompiliert ihn zu TorchScript
2. TorchScript-Tracing (Tracing)
2.1 Grundlegende Tracing-Methoden
Beispiel
import torch.nn as nn
# ── Modell definieren ──────────────────────────────────────
class SimpleNet(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
self.pool = nn.MaxPool2d(2)
self.fc = nn.Linear(128 * 8 * 8, 10)
def forward(self, x):
x = self.pool(torch.relu(self.conv1(x)))
x = self.pool(torch.relu(self.conv2(x)))
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
# Modellinstanz erstellen
model = SimpleNet()
model.eval()
# Beispiel-Eingabe
example_input = torch.randn(1, 3, 64, 64)
# ── Methode 1: torch.jit.trace-Tracing ───────────────────
# TorchScript durch Ausführen des Modells und Aufzeichnen der Operationen erstellen
traced_model = torch.jit.trace(model, example_input)
print("Verfolgtes Modell:")
print(traced_model)
# Modell speichern
traced_model.save("simple_net_traced.pt")
# Modell laden
loaded_model = torch.jit.load("simple_net_traced.pt")
# Inferenz mit dem geladenen Modell durchführen
output = loaded_model(example_input)
print(f"Ausgabeform: {output.shape}")
2.2 Einschränkungen des Tracing
Die Tracing-Methode hat einige Einschränkungen:
- Zeichnet nur Operationen auf, die tatsächlich ausgeführt werden
- Kontrollfluss (z. B. if, for) ist festgelegt
- Eingaben mit dynamischen Größen können Probleme verursachen
Beispiel
class DynamicModel(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x):
# Der Kontrollfluss ist während des Tracings festgelegt
if x.sum() > 0:
return x * 2
else:
return x / 2
model = DynamicModel()
model.eval()
# Während des Tracings wird der if-Zweig auf den während des Tracings genommenen Pfad festgelegt
example_input = torch.tensor([1.0])
traced = torch.jit.trace(model, example_input)
# Auch bei unterschiedlicher Eingabe wird derselbe Zweig ausgeführt
print(traced(torch.tensor([1.0]))) # 2
print(traced(torch.tensor([-1.0]))) # Das Ergebnis ist immer noch 2, nicht -0.5
Für Modelle mit Kontrollfluss sollten Sie TorchScript-Scripting anstelle von Tracing verwenden.
3. TorchScript-Scripting (Scripting)
3.1 Grundlegende Verwendung
Beispiel
# ── Verwendung des @torch.jit.script-Dekorators ───────────────
@torch.jit.script
def scripted_function(x: torch.Tensor) -> torch.Tensor:
"""Konvertiert Funktion mittels Scripting"""
if x.sum() > 0:
return x * 2
else:
return x / 2
# Testen der gescripteten Funktion
input1 = torch.tensor([1.0, 2.0])
input2 = torch.tensor([-1.0, -2.0])
print(scripted_function(input1)) # [2., 4.]
print(scripted_function(input2)) # [-0.5, -1.]
# ── Modell skripten ───────────────────────────────────
class ScriptableModel(nn.Module):
def __init__(self):
super().__init__()
self.fc = nn.Linear(10, 10)
@torch.jit.export
def forward(self, x: torch.Tensor) -> torch.Tensor:
return torch.relu(self.fc(x))
@torch.jit.export
def predict(self, x: torch.Tensor) -> torch.Tensor:
"""Zusätzliche Exportmethode"""
out = self.forward(x)
return torch.argmax(out, dim=1)
model = ScriptableModel()
# Modell skripten
scripted_model = torch.jit.script(model)
print(scripted_model)
# Speichern
scripted_model.save("scripted_model.pt")
3.2 Komplexe Modelle scripten
Beispiel
class ConditionModel(nn.Module):
def __init__(self, num_classes: int):
super().__init__()
self.num_classes = num_classes
self.features = nn.Sequential(
nn.Conv2d(3, 32, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1),
nn.ReLU(),
nn.AdaptiveAvgPool2d(1),
nn.Flatten()
)
self.classifier = nn.Linear(64, num_classes)
def forward(self, x: torch.Tensor, use_softmax: bool = False) -> torch.Tensor:
"""
Ein Modell, das dynamische Bedingungen unterstützt
"""
features = self.features(x)
logits = self.classifier(features)
if use_softmax:
return torch.softmax(logits, dim=1)
else:
return logits
def get_prediction(self, x: torch.Tensor) -> torch.Tensor:
"""Hilfsmethode"""
logits = self.forward(x, use_softmax=False)
return torch.argmax(logits, dim=1)
# Skripting
model = ConditionModel(num_classes=10)
scripted_model = torch.jit.script(model, example_inputs=(torch.randn(1, 3, 32, 32),))
# Test
test_input = torch.randn(2, 3, 32, 32)
output1 = scripted_model(test_input, use_softmax=False)
output2 = scripted_model(test_input, use_softmax=True)
print(f"Form der Logits-Ausgabe: {output1.shape}")
print(f"Form der Softmax-Ausgabe: {output2.shape}")
4. ONNX-Export
4.1 ONNX-Grundlagen
ONNX(Open Neural Network eXchange)ONNX ist ein offenes Austauschformat für neuronale Netze, das die Konvertierung von Modellen zwischen verschiedenen Deep-Learning-Frameworks unterstützt.
Vorteile von ONNX:
- Frameworkübergreifend: Unterstützt von PyTorch, TensorFlow, Caffe2 usw.
- Plattformübergreifend: Unterstützt CPU, GPU, Mobilgeräte und andere Plattformen
- Hardwareoptimierung: Kann ONNX Runtime für effiziente Inferenz nutzen
- Umfangreiche Tools: Eine Fülle von Optimierungswerkzeugen und Deployment-Lösungen
4.2 Grundlegender ONNX-Export
Beispiel
import torch.nn as nn
import torchvision
# ── Modell definieren ──────────────────────────────────────
class ImageClassifier(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 32, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1),
nn.ReLU(),
nn.AdaptiveAvgPool2d(1),
nn.Flatten()
)
self.classifier = nn.Linear(64, 10)
def forward(self, x):
x = self.features(x)
x = self.classifier(x)
return x
model = ImageClassifier()
model.eval()
# Beispiel-Eingabe
example_input = torch.randn(1, 3, 32, 32)
# ── Export nach ONNX ──────────────────────────────────
output_path = "image_classifier.onnx"
torch.onnx.export(
model,
example_input,
output_path,
export_params=True, # Modellparameter exportieren
opset_version=14, # ONNX-Version
do_constant_folding=True, # Konstantenfaltungsoptimierung
input_names=['input'], # Eingabe-Tensor-Namen
output_names=['output'], # Ausgabe-Tensor-Namen
dynamic_axes={
'input': {0: 'batch_size'}, # Dynamische Batch-Dimension
'output': {0: 'batch_size'}
}
)
print(f"Modell exportiert nach: {output_path}")
# Exportiertes Modell validieren
import onnx
onnx_model = onnx.load(output_path)
onnx.checker.check_model(onnx_model)
print("ONNX-Modellvalidierung bestanden!")
4.3 Komplexe Modelle exportieren
Beispiel
import torchvision.models as models
# Ein vortrainiertes Modell laden
model = models.resnet18(pretrained=True)
model.eval()
# Beispieleingabe (Standard-ResNet 224x224)
example_input = torch.randn(1, 3, 224, 224)
# Export
torch.onnx.export(
model,
example_input,
"resnet18.onnx",
export_params=True,
opset_version=14,
input_names=['input'],
output_names=['output'],
dynamic_axes={
'input': {0: 'batch_size', 2: 'height', 3: 'width'},
'output': {0: 'batch_size'}
}
)
print("ResNet18 wurde exportiert")
# Spezielle Layer während des Exports behandeln
# Für Layer, die eine spezielle Behandlung erfordern, einen Workaround verwenden
class ModelWithSpecialOps(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(3, 64, 3, padding=1)
def forward(self, x):
# F.interpolate anstelle des nn.functional-Alias verwenden
x = self.conv(x)
x = torch.nn.functional.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)
return x
model = ModelWithSpecialOps()
model.eval()
torch.onnx.export(
model,
torch.randn(1, 3, 32, 32),
"special_ops.onnx",
input_names=['input'],
output_names=['output'],
opset_version=14
)
5. Exportvalidierung und Optimierung
5.1 Validierung exportierter Modelle
Beispiel
import onnx
import onnxruntime as ort
def verify_onnx_model(onnx_path, pytorch_model, test_input):
"""
Überprüfen Sie die Ausgabekonsistenz zwischen dem ONNX-Modell und dem PyTorch-Modell
"""
# 1. ONNX-Modell prüfen
onnx_model = onnx.load(onnx_path)
onnx.checker.check_model(onnx_model)
print("✓ Validierung der ONNX-Modellstruktur bestanden")
# 2. PyTorch-Inferenz
pytorch_model.eval()
with torch.no_grad():
pytorch_output = pytorch_model(test_input)
# 3. ONNX-Runtime-Inferenz
ort_session = ort.InferenceSession(onnx_path)
ort_output = ort_session.run(None, {'input': test_input.numpy()})[0]
# 4. Ausgaben vergleichen
diff = np.abs(pytorch_output.numpy() - ort_output).max()
print(f"✓ Maximale Ausgabendifferenz: {diff:.6f}")
if diff < 1e-5:
print("✓ Ausgabevalidierung bestanden")
return True
else:
print("⚠ Ausgabendifferenz ist groß")
return False
# Verwendungsbeispiel
model = ImageClassifier()
model.eval()
# Zuerst exportieren
torch.onnx.export(
model,
torch.randn(1, 3, 32, 32),
"test_model.onnx",
input_names=['input'],
output_names=['output']
)
# Validieren
test_input = torch.randn(2, 3, 32, 32)
verify_onnx_model("test_model.onnx", model, test_input)
5.2 ONNX-Runtime-Optimierung
Beispiel
from onnxruntime import GraphOptimizationLevel
# Inferenzsitzung erstellen und Optimierung konfigurieren
def create_optimized_session(onnx_path, providers=None):
if providers is None:
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
# Sitzungsoptionen konfigurieren
sess_options = ort.SessionOptions()
# Graphoptimierung aktivieren
sess_options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL
# Weitere Optimierungen aktivieren
sess_options.intra_op_num_threads = 4
sess_options.inter_op_num_threads = 4
# Session erstellen
session = ort.InferenceSession(onnx_path, sess_options, providers=providers)
return session
# Die optimierte Session für die Inferenz verwenden
session = create_optimized_session("resnet18.onnx")
# Eingabe- und Ausgabenamen abrufen
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
# Inferenz
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
output = session.run([output_name], {input_name: input_data})[0]
print(f"Form der Inferenzausgabe: {output.shape}")
5.3 Modellquantisierung
Beispiel
import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType
def quantize_onnx_model(input_path, output_path):
"""
ONNX-Modell dynamisch quantisieren
Modellgröße reduzieren und Inferenz beschleunigen
"""
# Dynamische Quantisierung (keine Kalibrierungsdaten erforderlich)
quantize_dynamic(
input_path,
output_path,
# Zu quantisierende Gewichtstypen angeben
weight_type=QuantType.QInt8
)
# Dateigrößenvergleich abrufen
import os
original_size = os.path.getsize(input_path) / 1024 / 1024
quantized_size = os.path.getsize(output_path) / 1024 / 1024
print(f"Originalmodell: {original_size:.2f} MB")
print(f"Quantisiertes Modell: {quantized_size:.2f} MB")
print(f"Kompressionsverhältnis: {original_size/quantized_size:.2f}x")
# Verwendung
quantize_onnx_model("resnet18.onnx", "resnet18_quantized.onnx")
6. Mobile Bereitstellung
6.1 Export in mobile Formate
Beispiel
# 1. Das Modell tracen
model = ImageClassifier()
model.eval()
example_input = torch.randn(1, 3, 32, 32)
traced_model = torch.jit.trace(model, example_input)
# 2. Das Modell optimieren (mobile Optimierung)
optimized_model = torch.jit.optimize_for_inference(traced_model)
# 3. Mobiles Modell speichern
optimized_model.save("mobile_model.pt")
# Methode 2: torchmobile verwenden (falls verfügbar)
# torchmobile wird für ein leichteres mobiles Deployment verwendet
# Bitte beachten Sie die offizielle Dokumentation zur Cross-Kompilierung
# Methode 3: Nach ONNX exportieren und eine mobile Laufzeitumgebung verwenden
# iOS: Core ML Tools verwenden
# Android: ONNX Runtime Mobile verwenden
print("Mobiles Modell ist bereit")
6.2 Core ML-Export (iOS)
Für die iOS-Plattform können Sie Core ML Tools verwenden, um das Modell in das Core ML-Format zu konvertieren. Sie müssen es zuerst in das ONNX-Format konvertieren und dann Core ML Tools für die Konvertierung verwenden.
Beispiel
# pip install coremltools
# Schritte für den Core ML-Export
# Hinweis: Dies ist nur unter macOS verfügbar
# Schritt 1: Zuerst das PyTorch-Modell nach ONNX exportieren
# from pytorch_example import your_model
# model = your_model()
# model.eval()
# example_input = torch.randn(1, 3, 224, 224)
# torch.onnx.export(
# model,
# example_input,
# "model.onnx",
# input_names=['input'],
# output_names=['output']
# )
# Schritt 2: Core ML Tools verwenden, um in das Core ML-Format zu konvertieren
# from coremltools.converters import onnx as onnx_coreml
# coreml_model = onnx_coreml.convert(
# model="model.onnx",
# minimum_deployment_target='13',
# image_input_names=['input']
# )
# coreml_model.save("model.mlmodel")
# Verwendungsbeispiel
print("iOS-Bereitstellung erfordert eine macOS-Umgebung")
print("Detaillierte Schritte: 1. pip install coremltools 2. ONNX exportieren 3. Mit coremltools konvertieren")
7. Bewährte Verfahren und häufige Probleme
7.1 Fehlerbehebung beim Export
| Problem | Ursache | Lösung |
|---|---|---|
| Export fehlgeschlagen | Nicht unterstützte Operation | Verwenden Sie eine höhere opset_version, oder ersetzen Sie die Operation |
| Inkonsistente Ausgabe | Dynamischer Kontrollfluss | Verwenden Sie Scripting anstelle von Tracing, oder korrigieren Sie die Eingabegröße |
| ONNX-Runtime-Fehler | Operator nicht unterstützt | Überprüfen Sie die Liste der von ONNX unterstützten Operationen |
7.2 Vergleich der Exportformate
| Format | Vorteile | Nachteile | Anwendbare Szenarien |
|---|---|---|---|
| TorchScript (.pt) | Native PyTorch-Unterstützung | Geringe frameworkübergreifende Unterstützung | Desktop-/Server-Bereitstellung |
| ONNX (.onnx) | Frameworkübergreifend, plattformübergreifend | Einige Operationen werden nicht unterstützt | Allgemeine Bereitstellung |
7.3 Export-Checkliste
Beispiel
def export_checklist(model, example_input):
"""Prüfungen vor dem Export"""
model.eval()
# 1. Sicherstellen, dass sich das Modell im Inferenzmodus befindet
print(f"Modellmodus: {'eval' if not model.training else 'train'}")
# 2. Eingabegröße prüfen
print(f"Form der Beispieleingabe: {example_input.shape}")
# 3. Forward-Pass validieren
with torch.no_grad():
output = model(example_input)
print(f"Ausgabeform: {output.shape}")
# 4. Prüfen auf nicht serialisierbare Operationen
# z. B. Lambda-Funktionen
for name, module in model.named_modules():
if hasattr(module, '__call__'):
pass # Benutzerdefinierte Module prüfen
print("✓ Exportprüfung abgeschlossen")
# Verwendungsbeispiel
model = ImageClassifier()
example_input = torch.randn(1, 3, 32, 32)
export_checklist(model, example_input)
7.4 Vollständiger Exportprozess
Beispiel
class ProductionModel(nn.Module):
"""Produktionsreifes Modell"""
def __init__(self):
super().__init__()
self.backbone = nn.Sequential(
nn.Conv2d(3, 64, 7, stride=2, padding=3),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.MaxPool2d(3, stride=2, padding=1),
# ... weitere Schichten
)
self.head = nn.Linear(512, 10)
def forward(self, x, return_features=False):
features = self.backbone(x)
features = torch.nn.functional.adaptive_avg_pool2d(features, 1).flatten(1)
if return_features:
return features
return self.head(features)
# Vollständiger Export-Workflow
model = ProductionModel()
model.load_state_dict(torch.load("weights.pth")) # Gewichte laden
model.eval()
# Beispieleingabe vorbereiten
example_input = torch.randn(1, 3, 224, 224)
# 1. TorchScript tracen
traced = torch.jit.trace(model, example_input)
traced.save("model_traced.pt")
# 2. ONNX exportieren
torch.onnx.export(
traced,
example_input,
"model.onnx",
input_names=['input'],
output_names=['output'],
dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}},
opset_version=14
)
# 3. ONNX optimieren
# onnx-simplifier verwenden, um redundante Operationen zu entfernen
# onnxsim model.onnx model_simplified.onnx
print("Alle Formate erfolgreich exportiert!")