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
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

# Beispiel für Tracing-Einschränkungen

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

import torch

# ── 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

# Komplexeres Scripting-Beispiel: ein Modell mit Bedingungen

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
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

# Das vollständige ResNet-Modell exportieren

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 numpy as np
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("&#x26a0; 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

import onnxruntime as ort
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

# ONNX-Modellquantisierung

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

# Methode 1: TorchScript Mobile-Bereitstellung
# 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

# coremltools muss installiert werden
# 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

# Checkliste vor dem Export

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

# Vollständiges Beispiel für den Modell-Export-Workflow

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!")
Andere Erweiterungen