西川和久の不定期コラム

強力なGPUを搭載した「Jetson Nano」でAIを走らせてみた

 6月に掲載した記事「CUDAコア128基のGPUを搭載したAI/深層学習向け『Jetson Nano開発者キット』を試す」において、NVIDIAが3月中旬に発表したNVIDIA「Jetson Nano」開発者キットの紹介、そしてMNISTなどを実行させ簡単なベンチマークテストを行なった。後編ではもう少しソフトウェア的な部分を解説したい。

ベンチマークテストで使ったMNISTの中身は!?

 前編のベンチマークテストで使ったMNISTのコードは、PytorchがGitHubで公開しているサンプルの1つだ。--no-cudaのオプションがあるため、簡単にCUDA オン/オフの比較ができ、さらにWSLなどほかの環境でも作動するため、今回の用途にピッタリだった。内容を実際に見てみたところ、コードは以下のとおりだった。

 少し長いため1/2と2/2で2分割している。前半部分が、class Netでモデル定義、def trainは学習、def testはテスト……と、今回の中核をなすコードとなる。後半部分は、コマンドラインのArgument処理、MNISTデータの読み込み、def mainなど一般的な処理系だ。

 まず前半の最初の数行で使用するライブラリなどのインポート。以降は、Pythonの構文より、Pytorch固有のファンクションなどが多く、AIも含めこのあたりの知識がないと一見して何をしているのかさっぱりわからない内容となる(結局筆者も完全には理解できなかった)。

PytorchがGitHubで公開しているMNIST(1/2)

from __future__ import print_function
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import datetime
from torchvision import datasets, transforms

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.fc1 = nn.Linear(4*4*50, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2, 2)
        x = x.view(-1, 4*4*50)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    print('start:', datetime.datetime.today())
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))

def test(args, model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

    print('end:', datetime.datetime.today())

# ここでいったん分割(コピペして実行する場合は後半とつなげる)

 class Netにあるモデルの定義は、MNISTのピクセルサイズが28×28(後述の表示プログラム参照)なので1次元ベクトルで784次元。モデル自体はFeed Forwardネットワーク。入力層800(4×4×50)、中間層500、出力層10……などとなっている。確認のためclass Netの定義までを切り出し(return F.log_softmax(x, dim=1)まで)、コードの下へ

net = Net()
print(net)

を加え結果は、

Net(
  (conv1): Conv2d(1, 20, kernel_size=(5, 5), stride=(1, 1))
  (conv2): Conv2d(20, 50, kernel_size=(5, 5), stride=(1, 1))
  (fc1): Linear(in_features=800, out_features=500, bias=True)
  (fc2): Linear(in_features=500, out_features=10, bias=True)
)

となった。話はやや前後するがコード後半で

for epoch in range(1, args.epochs + 1):

という部分がある。これはコード前半で定義している学習/テストを実行する回数で、デフォルトは10となっている。しかしGUIのまま実行回数が増えるとメモリ不足になるため、

CUI起動設定
$ sudo systemctl set-default multi-user.target
GUIに戻す
$ sudo systemctl set-default graphical.target

といったようにCUI起動設定にして再起動するとメモリに余裕ができる。この状態でepoch数を50にして学習(train)とテスト(test)のロスの関係をグラフにしてみた(ただしJetson Nanoだと途中で熱暴走したため、WSL上で演算させた)。さすがにconsoleの表示をコピペしてグラフ化するのは面倒なので、結果をCSVで保存、Excelで読み込んでグラフ化している。PythonでCSVを書き出すには、

PythonでCSVを書き出す(抜粋)

# CSVライブラリ読込み
import csv

    for epoch in range(1, args.epochs + 1):
        # train()とtest()からlossを戻すように変更
        loss0 = train(args, model, device, train_loader, optimizer, epoch)
        loss1 = test(args, model, device, test_loader)
        # CSV書き出し
        with open('/mnt/c/WSL-Tools/train.csv', 'a') as f:
            writer = csv.writer(f)
            writer.writerow([epoch, loss0, loss1])

とすればいい。

 もともとのコードではtrain()とtest()の戻り値はないが、lossを戻すように変更。openの'a'は追記書込み。最後の行、writer.writerowでCSVに書き出す並びを決めている。ここでは“回数,学習ロス,テストロス”の並びでCSVファイルが作られる。

epoch数(横軸)と学習loss(青)、テストloss(橙)の関係

 グラフを見ると25回辺りから学習、テストのロス共にほぼ横ばいになっているのがわかる。さらに学習は回数が増えるとじょじょにではあるがロスが減っているものの、テストのロスは変わらないため、必要以上に回数を増やしても意味がないこともわかる。

 consoleに表示している数値を眺めていても、何となくロスが減ってるな……と思うが、こうしてグラフ化すると「なるほど確かに学んでいる!」と目に見えるのでおもしろい(笑)。

PytorchがGitHubで公開しているMNIST(2/2)

def main():
    # 1)ここでコマンドのArgumentを解釈
    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
    parser.add_argument('--batch-size', type=int, default=64, metavar='N',
                        help='input batch size for training (default: 64)')
    parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
                        help='input batch size for testing (default: 1000)')
    parser.add_argument('--epochs', type=int, default=10, metavar='N',
                        help='number of epochs to train (default: 10)')
    parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
                        help='learning rate (default: 0.01)')
    parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
                        help='SGD momentum (default: 0.5)')
    parser.add_argument('--no-cuda', action='store_true', default=False,
                        help='disables CUDA training')
    parser.add_argument('--seed', type=int, default=1, metavar='S',
                        help='random seed (default: 1)')
    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                        help='how many batches to wait before logging training status')

    parser.add_argument('--save-model', action='store_true', default=False,
                        help='For Saving the current Model')
    args = parser.parse_args()
    # 1)ここまで
    use_cuda = not args.no_cuda and torch.cuda.is_available()
    torch.manual_seed(args.seed)

    device = torch.device("cuda" if use_cuda else "cpu")

    kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}

    # 2)学習/テスト用のデータをダウンロード
    train_loader = torch.utils.data.DataLoader(
        datasets.MNIST('../data', train=True, download=True,
                       transform=transforms.Compose([
                           transforms.ToTensor(),
                           transforms.Normalize((0.1307,), (0.3081,))
                       ])),
        batch_size=args.batch_size, shuffle=True, **kwargs)

    test_loader = torch.utils.data.DataLoader(
        datasets.MNIST('../data', train=False,
                       transform=transforms.Compose([
                           transforms.ToTensor(),
                           transforms.Normalize((0.1307,), (0.3081,))
                       ])),
        batch_size=args.test_batch_size, shuffle=True, **kwargs)
    # 2)ここまで

    # 3) モデルの設定、学習/テスト
    model = Net().to(device)
    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)
    for epoch in range(1, args.epochs + 1):
        train(args, model, device, train_loader, optimizer, epoch)
        test(args, model, device, test_loader)
    # 3)ここまで
    if (args.save_model):
        torch.save(model.state_dict(),"mnist_cnn.pt")

if __name__ == '__main__':
    main()

 後半のソースコードの1)の部分では、コマンドラインのArgumentを解釈している。--batch-size、--test-batch-size、--epochs、--lr、--momentum、--no-cuda、--seed、--log-interval、--save-modelとあり、それぞれ指定がなかった時のdefultも設定されている。

 2)の部分は、学習用とテスト用のデータを設定する。datasets.MNISTの'../data'がデータを保存するパス。すでにそのフォルダにある場合は、ダウンロードしない。

 ここで気になるのはデータの中身。今回使用したdatasets.MNISTデータは、下記のプログラムを実行すると(前編の初期設定+matplotlibのインストール済であればJetson Nano上のGUIで作動)、画面キャプチャのような手書きデータが大量に入っている。これを学習/比較して解を得るのがこのプログラムの概要となる。

 またコメントアウトしているdatasets.FashionMNISTに切り替えると、名前のとおり、ファッション系のアイテムが入ったデータセットとなる。いずれにしてもお試し用の学習/テスト用データを作るのはたいへんなので、あらかじめこのようなかたちで用意されているのは、初心者にとってたいへんありがたい。

MNISTデータ表示プログラム

import torch
import matplotlib.pyplot as plt
from torchvision import datasets, transforms

# 最低限の初期値セット
use_cuda = True
kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}

# データセットをダウンロード
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('../data', train=True, download=True,
    #datasets.FashionMNIST('../data', train=True, download=True,
    transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])),
    batch_size=64, shuffle=True, **kwargs)

# 1つ目のデータ表示
img_data = train_loader.dataset[0]
plt.imshow(img_data[0].numpy().reshape(28, 28), cmap='gray')
plt.show()
手書き文字データセット
ファッションデータセット

 3)はモデルの設定と、学習(train)/テスト(test)のループとなる。回数は1からepochs(デフォルト10)オプションの値+1。時間がかかるので前編のベンチマークテストではここを--epochs=1として1回にした。

 ……と、筆者が理解できたのははここまで(笑)。肝心のモデル、学習、テストの内容がイマイチわかっていない。前編が6月11日公開だったので、約1月間、仕事の合間に、個人のブログなども含め、いろいろJetson関連のプログラムを動かし遊んでいたものの、そもそも仕掛け(理論)を理解できてないため、コピペと変わらないコードしか書けず、それを記事にもできないため、オリジナル的なものはあきらめることにした。

 できれば、筆者が用意したデータ(グラビア写真)なら山のようにある(笑)ため、学習させテストするようなサンプルを作りたかったのだが……。記事掲載が遅れたうえに試せずに申しわけない。

jetson-inferenceで遊ぶ

 本来、前編でご紹介すべき内容だが、SDカードに書き込んだ「Jetson Nano Developer Kit SD Card Image New JP 4.2」に含まれるSDKと、GitHubで同社が公開している「jetson-inference」で、パッと見てAIだとわかるものを試してみたい。

 makeが必要なのだが、SDKにほとんど入っているため、あれこれインストールしなくとも簡単に実行することができる。

$ sudo apt-get install git cmake
$ git clone https://github.com/dusty-nv/jetson-inference
$ cd jetson-inference
$ git submodule update --init
$ mkdir build
$ cd build
$ cmake ../
$ make
※ make中何度かパネルを表示するが、そのまま[OK]で問題ない
jetson-inferenceをmake中

 make後、jetson-inference/build/aarch64/binにいろいろなサンプル画像とコマンドができている。たとえば「imagenet-camera」はJetson NanoへRaspberry Pi カメラを付け動物などを認識できるコマンドだ。

 また同じコマンド名.pyになっているのはPython版だ。機能的にはバイナリ版と同じなので、もし中身を触りたい時にはPython版を使った方が簡単だろう。このとき、インタープリタになるものの、AI処理に時間がかかるため処理系の違いはほとんど無視できる。

 今回カメラは増設していないため、単独で動くものとして「imagenet-console」と「detectnet-console」を試用した。前者は画像が何かを当てるもの、後者はたとえば写真に写っている複数の犬などを認識できる。百聞は一見にしかず。まずimagenet-consoleから。使い方はこんな感じだ。

$ imagenet-console orange_0.jpg output_0.jpg
※1 imagenet-console 入力画像名 出力画像名
※2 orange_0.jpgはjetson-inference/build/aarch64/binにサンプルとしてある

 入力画像が何かを判断して入力画像と同じ出力画像の上にオーバーレイ表示する。いくつか試した結果は以下のとおり。

入力画像 / orange
出力画像 / 97.559%でorange
入力画像 / Jetson Nano
出力画像 / 42.358%でradio, wireless
入力画像 / スマートフォン
出力画像 / 61.377%でremote control, remote
入力画像 / ノートPC
出力画像 / 76.953%でnotebook, computer
入力画像 / 猫
出力画像 / 72.070%でEgyptian cat

 orangeはもともとサンプルで入っていたものでさすがに高確率で一致。ノートPCも正解だ。猫は周囲に余計なものが写り込んでいるものの、猫とただしく判別。背景がシンプルでなくても大丈夫なようだ。Jetson Nanoがラジオ、スマートフォンがリモコン……は、あたらからずも遠からずといった感じか(笑)。

 入力画像サイズが大きいと時間がかかるため、サンプルはすべて600×600pixに統一した。このサイズならほぼ瞬時で演算が終わる。imagenet-consoleの詳細はここにあるので興味のある人はご覧いただきたい。

 次はdetectnet-console。これは適当な写真の持ち合わせがなかったため、サンプルをそのまま使用している。結果は8匹中3匹を認識。写真の角度だったり大きさだったりと、ほかの5匹は難しかったのだろう。コマンドの最後にcoco-dogとあるが、これは対象となるオブジェクトの種類が入る。

$ detectnet-console dog_0.jpg output_5.jpg coco-dog
※ coco-dogは犬のオブジェクト指定
入力画像 / 写真に犬が8匹写っている
出力画像 / 認識したのは3匹

 coco-dogのほかに、coco-bottle、coco-chair、coco-airplane、pednet、multiped、facenetなどのオブジェクトも指定できる(つまり公開しているデータがある)。該当する被写体が複数写っている写真があれば試してみるといいだろう。詳細はここにあるので参考にしてほしい。

 いかがだろうか。前編では数字の羅列でしかなかったAIが、これらのコマンドによってビジュアル化され、本当にAIが動いていると実感できる。また今回試用したimagenet-consoleとdetectnet-consoleは、入力画像フォーマットがjpg/png/tga/bmpと一般的なものなので、手持ちの写真で楽しめるのもポイントが高い。

 さらに、にmakeで作られた「imagenet-camera」と「detectnet-camera」コマンドは、名前のとおり、Jetson Nanoに取り付けたカメラで撮影した画像を入力に使い、imagenet-console/detectnet-consoleと同じ処理ができ、よりロボットっぽくなる感じだ。手持ちで取り付け可能なカメラをお持ちの方は、ぜひ挑戦していただきたい。


 以上、前編・後編と2回に分けNVIDIA「Jetson Nano開発者キット」の話を書いた。筆者のAIに対する理解が浅く、後編も結局サンプルを動かした程度になってしまったが、もう少し勉強し機会があれば、また何かのときにリベンジしたい。いずれにしても本記事がきっかけで、AIやJetson(Nano)に興味を持っていただければ幸いだ。

 なお、今回のレビューにあたって菱洋エレクトロからサンプル機材の提供を受けたが、さらに読者プレゼントとして2台ご用意いただいたので、興味がある人はぜひプレゼントコーナーから応募してほしい。