surgeon pytorch
0.0.4
Eine Bibliothek zum Überprüfen und Extrahieren von Zwischenschichten von PyTorch-Modellen.
Es kommt häufig vor, dass wir Zwischenschichten von PyTorch-Modellen untersuchen möchten, ohne den Code zu ändern. Dies kann nützlich sein, um Aufmerksamkeitsmatrizen von Sprachmodellen zu erhalten, Ebeneneinbettungen zu visualisieren oder eine Verlustfunktion auf Zwischenschichten anzuwenden. Manchmal möchten wir Unterteile des Modells extrahieren und unabhängig voneinander ausführen, entweder um sie zu debuggen oder um sie separat zu trainieren. All dies kann mit Surgeon durchgeführt werden, ohne eine Zeile des Originalmodells zu ändern.
$ pip install surgeon-pytorch Bei einem PyTorch-Modell können wir alle Ebenen mit get_layers anzeigen:
import torch
import torch . nn as nn
from surgeon_pytorch import Inspect , get_layers
class SomeModel ( nn . Module ):
def __init__ ( self ):
super (). __init__ ()
self . layer1 = nn . Linear ( 5 , 3 )
self . layer2 = nn . Linear ( 3 , 2 )
self . layer3 = nn . Linear ( 2 , 1 )
def forward ( self , x ):
x1 = self . layer1 ( x )
x2 = self . layer2 ( x1 )
y = self . layer3 ( x2 )
return y
model = SomeModel ()
print ( get_layers ( model )) # ['layer1', 'layer2', 'layer3'] Dann können wir unser zu prüfendes model mit Inspect umschließen und bei jedem Vorwärtsaufruf des neuen Modells werden wir auch die bereitgestellten Layer-Ausgaben ausgeben (im zweiten Rückgabewert):
model_wrapped = Inspect ( model , layer = 'layer2' )
x = torch . rand ( 1 , 5 )
y , x2 = model_wrapped ( x )
print ( x2 ) # tensor([[-0.2726, 0.0910]], grad_fn=<AddmmBackward0>)Wir können eine Liste der Ebenen bereitstellen:
model_wrapped = Inspect ( model , layer = [ 'layer1' , 'layer2' ])
x = torch . rand ( 1 , 5 )
y , [ x1 , x2 ] = model_wrapped ( x )
print ( x1 ) # tensor([[ 0.1739, 0.3844, -0.4724]], grad_fn=<AddmmBackward0>)
print ( x2 ) # tensor([[-0.2238, 0.0107]], grad_fn=<AddmmBackward0>)Wir können ein Wörterbuch bereitstellen, um benannte Ausgaben zu erhalten:
model_wrapped = Inspect ( model , layer = { 'layer1' : 'x1' , 'layer2' : 'x2' })
x = torch . rand ( 1 , 5 )
y , layers = model_wrapped ( x )
print ( layers )
"""
{
'x1': tensor([[ 0.3707, 0.6584, -0.2970]], grad_fn=<AddmmBackward0>),
'x2': tensor([[-0.1953, -0.3408]], grad_fn=<AddmmBackward0>)
}
""" model = Inspect (
model : nn . Module ,
layer : Union [ str , Sequence [ str ], Dict [ str , str ]],
keep_output : bool = True ,
) Bei einem PyTorch-Modell können wir mit get_nodes alle Zwischenknoten des Diagramms anzeigen:
import torch
import torch . nn as nn
from surgeon_pytorch import Extract , get_nodes
class SomeModel ( nn . Module ):
def __init__ ( self ):
super (). __init__ ()
self . layer1 = nn . Linear ( 5 , 3 )
self . layer2 = nn . Linear ( 3 , 2 )
self . layer3 = nn . Linear ( 1 , 1 )
def forward ( self , x ):
x1 = torch . relu ( self . layer1 ( x ))
x2 = torch . sigmoid ( self . layer2 ( x1 ))
y = self . layer3 ( x2 ). tanh ()
return y
model = SomeModel ()
print ( get_nodes ( model )) # ['x', 'layer1', 'relu', 'layer2', 'sigmoid', 'layer3', 'tanh'] Dann können wir Ausgaben mit Extract extrahieren, wodurch ein neues Modell erstellt wird, das den angeforderten Ausgabeknoten zurückgibt:
model_ext = Extract ( model , node_out = 'sigmoid' )
x = torch . rand ( 1 , 5 )
sigmoid = model_ext ( x )
print ( sigmoid ) # tensor([[0.5570, 0.3652]], grad_fn=<SigmoidBackward0>)Wir können auch ein Modell mit neuen Eingabeknoten extrahieren:
model_ext = Extract ( model , node_in = 'layer1' , node_out = 'sigmoid' )
layer1 = torch . rand ( 1 , 3 )
sigmoid = model_ext ( layer1 )
print ( sigmoid ) # tensor([[0.5444, 0.3965]], grad_fn=<SigmoidBackward0>)Wir können auch mehrere Ein- und Ausgänge bereitstellen und diese benennen:
model_ext = Extract ( model , node_in = { 'layer1' : 'x' }, node_out = { 'sigmoid' : 'y1' , 'relu' : 'y2' })
out = model_ext ( x = torch . rand ( 1 , 3 ))
print ( out )
"""
{
'y1': tensor([[0.4437, 0.7152]], grad_fn=<SigmoidBackward0>),
'y2': tensor([[0.0555, 0.9014, 0.8297]]),
}
""" Beachten Sie, dass das Ändern eines Eingabeknotens möglicherweise nicht ausreicht, um das Diagramm auszuschneiden (es könnten andere Abhängigkeiten bestehen, die mit vorherigen Eingaben verbunden sind). Um alle Eingaben des neuen Diagramms anzuzeigen, können wir model_ext.summary aufrufen, was uns einen Überblick über alle erforderlichen Eingaben und zurückgegebenen Ausgaben gibt:
import torch
import torch . nn as nn
from surgeon_pytorch import Extract , get_nodes
class SomeModel ( nn . Module ):
def __init__ ( self ):
super (). __init__ ()
self . layer1a = nn . Linear ( 2 , 2 )
self . layer1b = nn . Linear ( 2 , 2 )
self . layer2 = nn . Linear ( 2 , 1 )
def forward ( self , x ):
a = self . layer1a ( x )
b = self . layer1b ( x )
c = torch . add ( a , b )
y = self . layer2 ( c )
return y
model = SomeModel ()
print ( get_nodes ( model )) # ['x', 'layer1a', 'layer1b', 'add', 'layer2']
model_ext = Extract ( model , node_in = { 'layer1a' : 'my_input' }, node_out = { 'add' : 'my_add' })
print ( model_ext . summary ) # {'input': ('x', 'my_input'), 'output': {'my_add': add}}
out = model_ext ( x = torch . rand ( 1 , 2 ), my_input = torch . rand ( 1 , 2 ))
print ( out ) # {'my_add': tensor([[ 0.3722, -0.6843]], grad_fn=<AddBackward0>)} model = Extract (
model : nn . Module ,
node_in : Optional [ Union [ str , Sequence [ str ], Dict [ str , str ]]] = None ,
node_out : Optional [ Union [ str , Sequence [ str ], Dict [ str , str ]]] = None ,
tracer : Optional [ Type [ Tracer ]] = None , # Tracer class used, default: torch.fx.Tracer
concrete_args : Optional [ Dict [ str , Any ]] = None , # Tracer concrete_args, default: None
keep_output : bool = None , # Set to `True` to return original outputs as first argument, default: True except if node_out are provided
share_modules : bool = False , # Set to true if you want to share module weights with original model
) Die Inspect -Klasse führt immer das gesamte als Eingabe bereitgestellte Modell aus und verwendet spezielle Hooks, um die Tensorwerte während ihres Durchflusses aufzuzeichnen. Dieser Ansatz hat den Vorteil, dass (1) wir kein neues Modul erstellen und (2) er einen dynamischen Ausführungsgraphen ermöglicht (z. B. for Schleifen und if -Anweisungen, die von Eingaben abhängen). Die Nachteile von Inspect bestehen darin, dass (1) einige Berechnungen verschwendet werden, wenn wir nur einen Teil des Modells ausführen müssen, und (2) wir nur Werte von nn.Module Ebenen ausgeben können – keine Zwischenfunktionswerte.
Die Extract -Klasse erstellt mithilfe symbolischer Ablaufverfolgung ein völlig neues Modell. Die Vorteile dieses Ansatzes bestehen darin, dass (1) wir das Diagramm an einer beliebigen Stelle zuschneiden können und ein neues Modell erhalten, das nur diesen Teil berechnet, (2) wir Werte aus Zwischenfunktionen (nicht nur Schichten) extrahieren können und (3) wir auch Änderungen vornehmen können Eingabetensoren. Der Nachteil von Extract besteht darin, dass nur statische Diagramme zulässig sind (beachten Sie, dass die meisten Modelle über statische Diagramme verfügen).
