surgeon pytorch
0.0.4
Une bibliothèque pour inspecter et extraire les couches intermédiaires des modèles PyTorch.
Il arrive souvent que nous souhaitions inspecter les couches intermédiaires des modèles PyTorch sans modifier le code. Cela peut être utile pour attirer l'attention sur les matrices des modèles de langage, visualiser les intégrations de couches ou appliquer une fonction de perte aux couches intermédiaires. Parfois, nous souhaitons extraire des sous-parties du modèle et les exécuter indépendamment, soit pour les déboguer, soit pour les entraîner séparément. Tout cela peut être fait avec Surgeon sans changer une seule ligne du modèle original.
$ pip install surgeon-pytorch Étant donné un modèle PyTorch, nous pouvons afficher toutes les couches en utilisant get_layers :
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'] Ensuite, nous pouvons envelopper notre model à inspecter à l'aide Inspect et, à chaque appel direct du nouveau modèle, nous afficherons également les sorties de couche fournies (dans la deuxième valeur de retour) :
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>)Nous pouvons fournir une liste de couches :
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>)Nous pouvons fournir un dictionnaire pour obtenir des sorties nommées :
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 ,
) Étant donné un modèle PyTorch, nous pouvons afficher tous les nœuds intermédiaires du graphique en utilisant get_nodes :
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'] Ensuite, nous pouvons extraire les sorties en utilisant Extract , ce qui créera un nouveau modèle qui renvoie le nœud de sortie demandé :
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>)Nous pouvons également extraire un modèle avec de nouveaux nœuds d'entrée :
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>)Nous pouvons également fournir plusieurs entrées et sorties et les nommer :
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]]),
}
""" Notez que changer un nœud d'entrée peut ne pas suffire à couper le graphique (il peut y avoir d'autres dépendances liées aux entrées précédentes). Pour afficher toutes les entrées du nouveau graphique, nous pouvons appeler model_ext.summary qui nous donnera un aperçu de toutes les entrées requises et des sorties renvoyées :
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
) La classe Inspect exécute toujours l'intégralité du modèle fourni en entrée et utilise des hooks spéciaux pour enregistrer les valeurs du tenseur au fur et à mesure de leur passage. Cette approche présente les avantages suivants : (1) nous ne créons pas de nouveau module (2) elle permet un graphe d'exécution dynamique (c'est-à-dire des boucles for et des instructions if qui dépendent des entrées). Les inconvénients d' Inspect sont que (1) si nous n'avons besoin d'exécuter qu'une partie du modèle, certains calculs sont inutiles, et (2) nous ne pouvons générer que des valeurs à partir des couches nn.Module – pas de valeurs de fonction intermédiaires.
La classe Extract crée un tout nouveau modèle à l'aide du traçage symbolique. Les avantages de cette approche sont (1) nous pouvons recadrer le graphique n'importe où et obtenir un nouveau modèle qui calcule uniquement cette partie, (2) nous pouvons extraire des valeurs de fonctions intermédiaires (pas seulement des couches) et (3) nous pouvons également modifier tenseurs d'entrée. L'inconvénient d' Extract est que seuls les graphiques statiques sont autorisés (notez que la plupart des modèles ont des graphiques statiques).
