PyTorch torch.inference_mode Funktion
PyTorch torch Referenzhandbuch
torch.inference_modeist ein Kontextmanager in PyTorch, der für den Inferenzmodus verwendet wird. Er ist strenger,torch.no_gradnicht nur die Gradientenberechnung deaktivierend, sondern auch alle Tracking-Funktionen der Autograd-Engine deaktivierend.
Dies ist während der Modellinferenz effizienterno_gradund kann den Speicherverbrauch weiter reduzieren und die Inferenzgeschwindigkeit verbessern.
Funktionsdefinition
torch.inference_mode(mode=True)
Parameters:
mode(bool, optional): WennTrue(Standardwert), aktiviert den Inferenzmodus; wennFalse, wird der Inferenzmodus beendet. Standardwert istTrue。
Rückgabewert:
- Gibt einen Kontextmanager zurück, in dem Gradienten und autograd deaktiviert sind.
Anwendungsbeispiele
Beispiel 1: Grundlegende Verwendung
Beispiel
import torch
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
# Im inference_mode-Kontext
with torch.inference_mode():
y = x * 2
print("In inference_mode:", y.requires_grad)
# Im no_grad-Kontext
with torch.no_grad():
z = x * 2
print("In no_grad:", z.requires_grad)
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
# Im inference_mode-Kontext
with torch.inference_mode():
y = x * 2
print("In inference_mode:", y.requires_grad)
# Im no_grad-Kontext
with torch.no_grad():
z = x * 2
print("In no_grad:", z.requires_grad)
Die Ausgabe ist:
在 inference_mode 中: False 在 no_grad 中: False
Beispiel 2: Vergleich von no_grad und inference_mode
Beispiel
import torch
# Tensor erstellen
x = torch.randn(100, 100)
# In inference_mode
with torch.inference_mode():
# Viele Berechnungen durchgeführt
for _ in range(10):
x = torch.mm(x, x)
# Selbst nach Abschluss der Berechnung können Tensoren im Kontext nicht für Backpropagation verwendet werden
result = x.sum()
# Prüfen, ob er in einen Tensor umgewandelt werden kann, der Gradienten erfordert
print("Im inference_mode:", result.is_leaf)
# Dieselbe Berechnung in no_grad durchführen
x2 = torch.randn(100, 100)
with torch.no_grad():
for _ in range(10):
x2 = torch.mm(x2, x2)
result2 = x2.sum()
print("In no_grad:", result2.is_leaf)
# Tensor erstellen
x = torch.randn(100, 100)
# In inference_mode
with torch.inference_mode():
# Viele Berechnungen durchgeführt
for _ in range(10):
x = torch.mm(x, x)
# Selbst nach Abschluss der Berechnung können Tensoren im Kontext nicht für Backpropagation verwendet werden
result = x.sum()
# Prüfen, ob er in einen Tensor umgewandelt werden kann, der Gradienten erfordert
print("Im inference_mode:", result.is_leaf)
# Dieselbe Berechnung in no_grad durchführen
x2 = torch.randn(100, 100)
with torch.no_grad():
for _ in range(10):
x2 = torch.mm(x2, x2)
result2 = x2.sum()
print("In no_grad:", result2.is_leaf)
Die Ausgabe ist:
在 inference_mode 中: False 在 no_grad 中: True
Beispiel 3: Modellinferenz
Beispiel
import torch
import torch.nn as nn
# Ein einfaches Modell definieren
model = nn.Sequential(
nn.Linear(10, 20),
nn.ReLU(),
nn.Linear(20, 5)
)
model.eval()
# Eingabedaten erstellen
x = torch.randn(100, 10)
# inference_mode für die Inferenz verwenden
with torch.inference_mode():
output = model(x)
print("Ausgabeform:", output.shape)
print("Ausgabe requires_grad:", output.requires_grad)
# Kann auch der Dekorator verwendet werden
@torch.inference_mode()
def predict(x):
return model(x)
result = predict(x)
print("Dekorator-Methode - Ausgabeform:", result.shape)
import torch.nn as nn
# Ein einfaches Modell definieren
model = nn.Sequential(
nn.Linear(10, 20),
nn.ReLU(),
nn.Linear(20, 5)
)
model.eval()
# Eingabedaten erstellen
x = torch.randn(100, 10)
# inference_mode für die Inferenz verwenden
with torch.inference_mode():
output = model(x)
print("Ausgabeform:", output.shape)
print("Ausgabe requires_grad:", output.requires_grad)
# Kann auch der Dekorator verwendet werden
@torch.inference_mode()
def predict(x):
return model(x)
result = predict(x)
print("Dekorator-Methode - Ausgabeform:", result.shape)
Die Ausgabe ist:
输出形状: torch.Size([100, 5]) 输出 requires_grad: False 装饰器方式 - 输出形状: torch.Size([100, 5])
Beispiel 4: Vergleich der Speicheroptimierung
Beispiel
import torch
import torch.nn as nn
model = nn.Sequential(
nn.Linear(1000, 1000),
nn.ReLU(),
nn.Linear(1000, 1000),
nn.ReLU(),
nn.Linear(1000, 10)
)
# Testen der Speichernutzung verschiedener Modi
x = torch.randn(50, 1000)
print("Ohne Kontextmanager:")
_ = model(x)
print("\nVerwendung von no_grad:")
with torch.no_grad():
_ = model(x)
print("\nVerwendung von inference_mode:")
with torch.inference_mode():
_ = model(x)
import torch.nn as nn
model = nn.Sequential(
nn.Linear(1000, 1000),
nn.ReLU(),
nn.Linear(1000, 1000),
nn.ReLU(),
nn.Linear(1000, 10)
)
# Testen der Speichernutzung verschiedener Modi
x = torch.randn(50, 1000)
print("Ohne Kontextmanager:")
_ = model(x)
print("\nVerwendung von no_grad:")
with torch.no_grad():
_ = model(x)
print("\nVerwendung von inference_mode:")
with torch.inference_mode():
_ = model(x)
Verwendunginference_modekann den Speicher weiter optimieren, da es die Autograd-Engine vollständig deaktiviert.
Verwandte Funktionen
torch.no_grad(): Deaktiviert die Gradientenberechnung, behält jedoch einen Teil der Autograd-Funktionalität bei.torch.enable_grad(): Aktiviert die Gradientenberechnung.torch.is_inference_mode_enabled(): Prüft, ob der Inferenzmodus aktiviert ist.
Hinweise
inference_modeVergleichenno_gradstrikter, deaktiviert mehr Funktionalität.- In
inference_modeDarin erstellte Tensoren werden als Nicht-Blatt-Knoten markiert und können nicht für die Backpropagation verwendet werden. - Es wird empfohlen, es während der Modellinferenz und -evaluierung zu verwenden
inference_modeum die beste Leistung zu erzielen. inference_modeKann nichtno_gradverschachtelt werden.
Andere Erweiterungen