surgeon pytorch
0.0.4
Библиотека для проверки и извлечения промежуточных слоев моделей PyTorch.
Часто нам нужно проверить промежуточные слои моделей PyTorch без изменения кода. Это может быть полезно для получения матриц внимания языковых моделей, визуализации вложений слоев или применения функции потерь к промежуточным слоям. Иногда нам нужно извлечь части модели и запустить их независимо, либо для отладки, либо для обучения отдельно. Все это можно сделать с помощью Surgeon, не меняя ни одной строчки исходной модели.
$ pip install surgeon-pytorch Учитывая модель PyTorch, мы можем отобразить все слои, используя 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'] Затем мы можем обернуть нашу model для проверки с помощью Inspect , и при каждом прямом вызове новой модели мы также будем выводить предоставленные выходные данные слоя (во втором возвращаемом значении):
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>)Мы можем предоставить список слоев:
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>)Мы можем предоставить словарь для получения именованных выходных данных:
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 ,
) Учитывая модель PyTorch, мы можем отобразить все промежуточные узлы графа, используя 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'] Затем мы можем извлечь выходные данные с помощью Extract , который создаст новую модель, возвращающую запрошенный выходной узел:
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>)Мы также можем извлечь модель с новыми входными узлами:
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>)Мы также можем предоставить несколько входов и выходов и назвать их:
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]]),
}
""" Обратите внимание, что изменения входного узла может быть недостаточно для разрезания графа (могут существовать другие зависимости, связанные с предыдущими входными данными). Чтобы просмотреть все входные данные нового графика, мы можем вызвать model_ext.summary , который даст нам обзор всех необходимых входных данных и возвращаемых выходных данных:
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
) Класс Inspect всегда выполняет всю модель, предоставленную в качестве входных данных, и использует специальные перехватчики для записи значений тензора по мере их прохождения. Этот подход имеет преимущества: (1) мы не создаем новый модуль (2) он допускает динамический граф выполнения (т. е. циклы for и операторы if , которые зависят от входных данных). Недостатки Inspect заключаются в том, что (1) если нам нужно выполнить только часть модели, некоторые вычисления будут потрачены впустую, и (2) мы можем выводить только значения из слоев nn.Module – без промежуточных значений функции.
Класс Extract создает совершенно новую модель, используя символьную трассировку. Преимущества этого подхода: (1) мы можем обрезать график где угодно и получить новую модель, которая вычисляет только эту часть, (2) мы можем извлекать значения из промежуточных функций (не только слоев) и (3) мы также можем изменять входные тензоры. Недостатком Extract является то, что разрешены только статические графики (обратите внимание, что большинство моделей имеют статические графики).
