こすたろーんエンジニアの試行錯誤部屋

作成物の備忘録を書いていきますー

pytorchのdashboardのtensorboradXをinstallして試してみた

tensorboardXは、データフローグラフや学習関係の変数(lossやaccuracyなど)の可視化できるダッシュボードツールです
(tensorflowで使われるtensorboardのpytorch版)
今回以前に作成したsimsiamコードにtensorboardXを導入します
technoxs-stacker.hatenablog.com

目次

スポンサーリンク

この記事でわかること

・tensorboardXのインストール方法
・tensorboardXの基本的な使い方

1.実行環境

Jetson Xavier NX
ubuntu18.04
docker
python3.x
pytorch
->Jetson Xavier NX上におけるpytrorch環境構築は以下でやってますので、ご参考までに~ technoxs-stacker.hatenablog.com

2.インストール方法

pipでインストール可能です

pip install tensorboard
pip install tensorboardX

3.基本的な使い方

以下の流れでtensorboardXによる可視化を行うことができます
3.1必要なモジュールのインポート
3.2writerの定義
3.3数値の書き込み
3.4可視化ツールの立ち上げ

3.1必要なモジュールのインポート

必要なモジュールをimportします

# ------------------------------- 3.1
from tensorboardX import SummaryWriter
# -------------------------------

3.2writerの定義

書き込み用writerモジュールを定義します
第一引数logdirに出力結果を保存するディレクトリを指定します

def main_worker(gpu):
  ・
  ・
  ・
    # ------------------------------- 3.2
    # define writer
    writer = SummaryWriter(log_dir)
    # ------------------------------- 

    for epoch in range(start_epoch, epochs):
        adjust_learning_rate(optimizer, init_lr, epoch, epochs)

        # train for one epoch
        # ------------------------------- 3.2 -> 引数にwriterを追加
        train(train_loader, model, criterion, optimizer, epoch, gpu, print_freq, writer)
        # -------------------------------

        save_checkpoint({
            'epoch': epoch + 1,
            'arch': arch,
            'state_dict': model.state_dict(),
            'optimizer' : optimizer.state_dict(),
        }, is_best=False, filename='checkpoint_{:04d}.pth.tar'.format(epoch))

    torch.save(model.state_dict(),
                checkpoint_dir / 'latest.pth')

3.3 数値の書き込み

可視化したい変数をwriterに追加します
今回はloss値を可視化します
使用する関数はadd_scalarです

項目 関数
add_scalar writer.add_scalar("label", value, horizontal axis value)

train関数にadd_scalarを追記します

# ------------------------------- 3.3 -> 引数にwriterを追加
def train(train_loader, model, criterion, optimizer, epoch, gpu, print_freq, writer):
# ------------------------------- 
    batch_time = AverageMeter('Time', ':6.3f')
    data_time = AverageMeter('Data', ':6.3f')
    losses = AverageMeter('Loss', ':.4f')
    progress = ProgressMeter(
        len(train_loader),
        [batch_time, data_time, losses],
        prefix="Epoch: [{}]".format(epoch))

    # switch to train mode
    model.train()
    end = time.time()
    for i, (images, _) in enumerate(train_loader, start=epoch * len(train_loader)):
        # measure data loading time
        data_time.update(time.time() - end)
        images[0] = images[0].cuda(gpu, non_blocking=True)
        images[1] = images[1].cuda(gpu, non_blocking=True)

        # compute output and loss
        p1, p2, z1, z2 = model(x1=images[0], x2=images[1])

        # compute output and loss
        loss = -(criterion(p1, z2).mean() + criterion(p2, z1).mean()) * 0.5
        losses.update(loss.item(), images[0].size(0))

        # compute gradient and do SGD step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if i % print_freq == 0:
            progress.display(i)
        
        # ------------------------------- 3.3
        writer.add_scalar("train_loss", loss.item(), i)
        # -------------------------------

3.4 可視化ツールの立ち上げ

以下コマンドでダッシュボードツールが立ち上がります

tensorboard --logdir="path_to_log" --port=port_no

感想

学習が進んでいるか簡単に確認できます~
ヒストグラムも可視化できるのでいろいろ試してみたいですね

スポンサーリンク