機械学習モデルの解析: surgeon-pytorch

surgeon-pytorchとは

surgeon-pytorchとは機械学習モデルの層の解析のためのライブラリです。

surgeon-pytorchを使うと解析が楽に実装できるのでオススメです。

https://github.com/archinetai/surgeon-pytorch

主にできることは以下です。

• 層の名前の取得(get_layers)
• 計算ノードの名前の取得(get_nodes)
• 隠れ層の出力の取得(Inspect,Extract)
• 隠れ層の内部(計算ノード)の計算の取得(Extract)

surgeon-pytorchの使い方

モデルの情報を取得

import torch
import torch.nn as nn
from surgeon_pytorch import get_layers, 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(2, 1)

    def forward(self, x):
        x1 = self.layer1(x)
        x2 = self.layer2(x1)
        y = self.layer3(x2)
        return y
	
get_layers(model) #modelの層名を取得
get_nodes(model) #modelの計算ノード名を取得

モデルの隠れ層の出力を取得(Inspect)

from surgeon_pytorch import Inspect
model_wrapped = Inspect(model, layer='{layer名}')
y, hidden_output = model_wrapped(x)

Inspectでは、モデルの中間層の出力を出力するモデルを作成できる。

Inspectで作成したモデルで順伝播の計算を行った場合、モデルは出力層まで計算を行う。

したがって隠れ層の出力(hidden_out)のみ欲しい場合は、計算が無駄になってしまう。

隠れ層の出力(hidden_out)のみ欲しい場合は、モデルを新たに構築できるExtractを使うと良い。

API

model = Inspect(
    model: nn.Module,
    layer: Union[str, Sequence[str], Dict[str, str]],
    keep_output: bool = True,
)

モデルの隠れ層(の計算ノード)の出力を取得(Extract)

from surgeon_pytorch import Extract
model_ext = Inspect(model, node_out='{node名}')
hidden_output = model_ext(x)

Extractでは、モデルを新たに作成できる。

計算ノード単位で入力(node_inで定義)及び出力(node_outで定義)を定義できる。

API

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
)