PyProfでPyTorchのGPU上での動作をプロファイルする

DLProf以外の方法でもGPUの動作を確認できるプロファイリングできる方法があるぞ

より詳細を把握したいときに便利そうですね

DLProfのDeepLearningのプロファイリング

前回のブログではDLProfでプロファイリングしてTensorboardで確認しました。

Pyprofのプロファイリング

今回は別の方法でプロファイルする方法を確認してみます。pyprofだとDLProfと異なり、レイヤーごとの時間が分かるので動作がクリティカルに重いレイヤーを把握するのに役立ちます。

https://github.com/NVIDIA/PyProf

Quick startに沿って試します。

https://docs.nvidia.com/deeplearning/frameworks/pyprof-user-guide/quickstart.html

環境構築は下記の記事と同様です。

実行するコードはシンプルなコードlenet.pyを使用します。

必要なライブラリをインポートします。

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.cuda.profiler as profiler
import torch.optim as optim
try:
    from apex import amp
except ImportError:
    raise ImportError("Please install apex from https://www.github.com/nvidia/apex to run this example.")
from tqdm import tqdm
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--amp", help="setting amp option",
                    action="store_true")
args = parser.parse_args()

PyProfでプロファイリングできるように設定します。

import pyprof
pyprof.init()

lenetモデルクラスを定義します。

class LeNet5(nn.Module):

    def __init__(self):
        super(LeNet5, self).__init__()
        # 1 input image channel, 6 output channels, 5x5 square convolution
        # kernel
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        # an affine operation: y = Wx + b
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # Max pooling over a (2, 2) window
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        # If the size is a square you can only specify a single number
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, self.num_flat_features(x))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]  # all dimensions except the batch dimension
        num_features = 1
        for s in size:
            num_features *= s
        return num_features

autograd処理の際にプロファイルできるようにコンテキストマネージャーを作成しています。

https://pytorch.org/docs/stable/autograd.html#torch.autograd.profiler.emit_nvtx

Tensorコアが動作しているかどうかも確認したかったのでapexを使用しています。

https://github.com/NVIDIA/apex

with torch.autograd.profiler.emit_nvtx():

    net = LeNet5().cuda()

    input = torch.randn(1, 1, 32, 32).cuda()
    out = net(input)

    target = torch.randn(10)  # a dummy target, for example
    target = target.view(1, -1).cuda()  # make it the same shape as output
    criterion = nn.MSELoss()

    # create your optimizer
    optimizer = optim.SGD(net.parameters(), lr=0.01)
    if args.amp:
        opt_level = 'O1'
        net, optimizer = amp.initialize(net, optimizer, opt_level=opt_level)

    for i in tqdm(range(10)):
        # in your training loop:
        optimizer.zero_grad()  # zero the gradient buffers

        profiler.start()
        output = net(input)
        loss = criterion(output, target)
        if args.amp:
            with amp.scale_loss(loss, optimizer) as scaled_loss:
                scaled_loss.backward()
        else:
            loss.backward()
        optimizer.step()  # Does the update
    profiler.stop()

プロファイリング

下記のコマンドで実行します。

nsys profile -f true -o net --export sqlite python lenet.py

Nsight Systemで可視化をします。

下記からダウンロードしてインストールします。

https://developer.nvidia.com/nsight-systems#Platforms

インストールしてディレクトリに移動してNsight Systemを実行します。

cd ~/nsight-systems-2020.3.1
./bin/nsight-sys

実行すると下記のようになります。

プロファイリングデータ(net.qdrep)が作成されているのでそれをNsight Systemで開きます。

下記のように各オペレーションごとの時間がみれます。このままでは細かい情報がわからないので拡大します。

NVTXの行を拡大して確認します。NVTXはタイムラインにマーカーをつけてくれるのでタイムラインとレイヤーの対応関係が把握しやすくなります。

https://docs.nvidia.com/cuda/profiler-users-guide/index.html#nvtx

拡大するとLossの実行時間や各種レイヤーのバックワード時間を把握できます。

PyProfを使用するとDLProfとは違う観点でプロファイリングできるぞ!

Close Bitnami banner
Bitnami