PyTorchによるImageNet画像分類スクリプトの作り方

この記事は Deep Learning エンジニアの Dominic Monn （@dqmonn）氏が TECH x GAME COLLEGE のために寄稿していただいたものをQiita用にリライトしたものです。

その名前から期待できる通り、ImageNetというのはニューラルネットワークのことではありません。ImageNetは大規模な画像データベースであり、2万を超える分類から構成された、14億を超える画像を保有しています。

最近、ImageNetは訓練用データ・セットとして使用されることが多いのですが、新しく作られたニューラルネットワークの性能を評価する目的でも使われています。ニューラルネットワークを比較するために、データベース上で1,000近くの分類が使用されます。最先端のニューラルネットワークは通常90パーセントを超える精度を達成しています。

この記事で使用するソースコードについて

この記事で使用されている全コードは https://github.com/tech-x-college/pytorch-imagenet に公開されています。

前提条件

最先端の画像分類スクリプトを構築するにはある程度の処理能力が必要なので、ハイエンドのGPU環境で作業をするか、もしくは、Google Colaboratoryのような無料クラウドトレーニングプロパイダを使って環境を構築するようにしてください。
PyTorchはhttps://pytorch.org/get-started/locally/よりダウンロードできます。可能であればCUDAサポート版をインストールしてください。

データの準備

ImageNetは14億もの画像データを持ち、合計サイズが100GBを超える大規模なデータ・セットです。長期スパンで見ればそのサイズのデータセットをダウンロードするのは価値あることかもしれませんが、このクイックチュートリアルには明らかに必要のない量です。

この研究のため、スタンフォード大学は「Tiny Imagenet」という名前で、200の分類で、その1つの分類に対して500枚の訓練画像と100のテスト・検証画用像を公開するこを決めました。このデータセットでは最先端のパフォーマンスに到達することは不可能ですが、簡単な実験環境を構築することがずっと楽になります。

このチュートリアルに関しては、以下のリンクからTiny Imagenetをダウンロードしてください。(https://tiny-imagenet.herokuapp.com/)

ImageFolder と TrainLoader

大規模な画像フォルダを読み込むためには、'DataLoader'クラスを使い継続的に訓練用の新しい画像を読み込ませる前に、まず'torch.datasets'の'ImageFolder'クラスを用いたデータパイプラインを設定することが大切です。

全てのデータ・セットをただ直接読み込むのではなく、これらのクラスを使う理由はメモリの制限があるからです。とりわけ、ImageNet全てのデータ・セットを読み込みたい場合、メモリ領域はあっと言う間になくなってしまい、全てのシステムがフリーズしてしまいます。'ImageFolder'クラスは継続的かつ並列的に新しい画像を読み込み、処理し、そして'DataLoader'クラスにそれらの画像を渡します。そうすることで初めて、それらの画像をサンプリングし、バッチ処理として扱い、訓練回路に利用することができるようになります。

ImageFolder

ImageFolderクラスを使用するには、データが置かれている場所を指定する必要があります。最も簡単な方法は、スクリプトを呼び出す際に実行引数を指定することです。pythonファイルに次の行を追加することで、以下の'python3 script.py'を実行することができます。

parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
parser.add_argument('data', metavar='DIR',
                    help='path to dataset')

その後、このパス情報を使い、'ImageFolder'クラスを初期化することができます。データの正規化と水増しする方法に関しては注意が必要です。まず第一に、一定の結果を得るためにノイズを減らし、計算の複雑性を軽減します。その次に、合成によってデータ・セットを拡大させます。画像を反転させたりランダムに切り出すことによって、新たな課題と画像をニューラルネットワークに取り込みます。そうすることで、とくに少ないサイズのデータ・セットにおける過学習を減らすことができます。

traindir = os.path.join(args.data, 'train')     # /train/ を指定されたパスに追加
testdir = os.path.join(args.data, 'test')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])     # 正規化定数

train_dataset = datasets.ImageFolder(
        traindir,
        transforms.Compose([
            transforms.RandomResizedCrop(224),      # 画像をサイズ224に切り出しもしくはリサイズ
            transforms.RandomHorizontalFlip(),      # ランダムに画像をフリップ（水増し）
            transforms.ToTensor(),
            normalize,
        ]))

ここまできたら、同様の作業をテストセットにも行います。ただしテストセットには、'data augmentation' (データの水増し) は利用しません。

DataLoader

画像ローダーの準備が完了したら、データ読み込みパイプラインをとても簡単に構築することができます。適度なバッチサイズとデータをシャッフルすることと処理を4つのスレッドで並列処理することを指定します。それだけが必要な作業です。

    train_loader = torch.utils.data.DataLoader(
            train_dataset, batch_size=64, shuffle=True,
            num_workers=4, pin_memory=True)

Training Loop

次にすることは、訓練回路を作ることです。訓練回路は、構築したネットワークにデータを渡し、損失と精度を計算し、さらにそのモデルを最適化する役目を担っています。

モデルの設定

モデルの設定はかつてないほど容易になりました。このチュートリアルではAlexNetを使うことにします。AlexNetはImageNetの研究に使うため、6年ほど前にリリースされた歴史あるモデルの1つです。リリース当時、それは最先端のモデルであり、Top-5エラー率15.3パーセントの数値でした。今日ではさらに強力な計算能力を有したモデルがTop-5エラー率5パーセント以下を達成しています。

AlexNetは次のように設定します。

model = models.alexnet()

先に述べたように、AlexNetは大規模なモデルです。ベストパフォーマンスを引き出すためには、Cudaを使用可能な処理能力の高いマシン上で動かす必要があります。

if args.gpu is not None:
    model = model.cuda(args.gpu)

これで全ての準備は整いました。

精度についてのひとこと

ImageNetの研究は大抵の場合、Top-1パフォーマンスとTop-5パフォーマンスの性能で測定されます。ImageNetはとても多くの分類を有しているので、Top-1パフォーマンスは最高出力の精度(例：フルーツは'フルーツ'として分類されなければいけない)、Top-5パフォーマンスはトップ5の出力の精度(例：フルーツの場合、上位5つの出力のいずれかに'フルーツ'が含まれていなければいけない)を測定します。

最適化の設定

損失関数と最適化の設定も同様に簡単です。訓練の準備は、たった二行のコードでおしまいです。

model = models.alexnet()
criterion = nn.CrossEntropyLoss().cuda(args.gpu)
optimizer = torch.optim.SGD(model.parameters(), 0.01,
                                momentum=0.9,
                                weight_decay=1e-4)

始めの行では、損失関数を設定しています。この場合、交差エントロピー損失を用いています。

'model.parameters()'では、その最適化ルーチンに全ての訓練可能な変数を含めることを推奨しています。次のパラメターは学習率、そして、運動量と重み減衰についての値が続きます。

損失と精度の計算

ここまできてようやく訓練回路を実装することができます。複数のエポックに対して訓練を実行し、その都度、完全訓練パスをチェックし、検証実行が続きます。

その訓練パスに対して、以下の作業をする新しい関数を定義します。

データの一部を取り出す
そのデータをモデルに送り込む
損失を計算する
誤差逆伝播法を実行する
最適化ステップを実行する
損失と精度を記録する

これは一見とても長いリストに見えますが、PyTorchを使えば一行のコードでこれら全ての問題を解決することができます。

** データの一部取り出し **

データ読み取りに関してはすでに書き終えているので、下記のようなループを定義することのみ必要です。

for i, (input, target) in enumerate(train_loader):

モデルへの送り込み

こんなに易しいことはありません。ただモデル関数を呼び出すだけでいいのです。

output = model(input)

損失の計算

すでに損失関数を設定しているので、最後に残っている必要な作業は出力とラベルを比べることです。

loss = criterion(output, target)

**バックポップの実行とその最適化

これに関しても、Pytorch訓練の一部であるいくつかのコマンドを実行するだけです。まず、勾配知をリセットし、最適化ステップを追加する前に誤差逆伝播法を実行します。

optimizer.zero_grad()
loss.backward()
optimizer.step()

損失と精度の記録

最後に、解析のため、訓練の精度を計算し損失関数を記録します。そうすることで、これらの情報を出力して訓練の進歩を確認することができます。本来であれば、それらの情報をきれいなグラフにプロットするとよいのですが、今回のチュートリアルでは単純さを追求しているためスキップします。

以前言及したように、Top-1とTop-5の精度を計算することにしましょう。幸運にも、PyTorchはすぐに使えるサンプルコードをチュートリアル内で提供してくれています。つまり、自分たちでそのコードを書く必要はなく、単純にその関数を借用し、精度を配列に記録することができるのです。

検証

検証を実行するためのプロセスは、誤差逆伝播法と最適化以外は全く同じです。訓練データを使う代わりに、以前設定したテストデータを利用します。この他に必要な作業はもうなにもありません！

結果

訓練させてみましょう。最初に気がつくことは訓練に要する時間がMNISTタスクでの例に比べて多いということです。今回のケースでは単に大きなデータセットだけではなく、はるかに大規模なモデルを使用しています。これには、大きなメモリサイズと、優れたハードウェアが必要です。

100のエポック(これは数時間を要します！)を訓練してみると、Top-1精度52.9パーセント及びTop-5精度76.5パーセントという値に達しました。

次のステップ

AlexNetの代わりにさらに大規模で強力なモデルを利用する。
フルサイズImageNetのように、より大きなImageNetデータサイズを利用する。
'転移学習' を利用し、訓練時間を短くする。

結論

このチュートリアルでは、PyTorchが提供するデータローダーを利用して打規模なデータサイズを訓練させる方法を学習しました。また、あらかじめ構築されたネットワークを使って訓練をさせる方法、そして最終的には複雑なタスクを設定し、それを解決することができました。このモデルのパフォーマンスをさらに高める高度なテクニックはまだまだたくさんありますが、それはまた別の機会に取っておきましょう！

PyTorch による ImageNet 画像分類スクリプトの作り方