PyTorchでAlexNetをゼロから書く

紹介

この投稿は、PyTorchで最も人気のある畳み込みニューラルネットワークをゼロから構築するシリーズの続きです。こちらで前回の投稿を見ることができ、LeNet5を構築しました。この投稿では、コンピュータビジョンにおける最も重要なブレークスルーアルゴリズムの一つであるAlexNetを構築します。

まず、AlexNetのアーキテクチャを調査し理解することから始めます。その後、データセットCIFAR-10を読み込み、データにいくつかの前処理を行った後、すぐにコードに入ります。次に、PyTorchを使用してゼロからAlexNetを構築し、前処理されたデータでトレーニングを行います。最後に、トレーニングされたモデルを評価目的で未見の（テスト）データでテストします。

前提条件

ニューラルネットワークの知識は、この記事を理解するのに役立ちます。これには、ニューラルネットワークの異なる層（入力層、隠れ層、出力層）、活性化関数、最適化アルゴリズム（勾配降下法のバリエーション）、損失関数などに慣れていることが含まれます。さらに、この記事で提示されているコードスニペットを理解するためには、Pythonの文法とPyTorchライブラリに精通していることが不可欠です。

CNNの理解は不可欠です。これには、畳み込み層、プーリング層、そしてそれらが入力データから特徴を抽出する役割に関する知識が含まれます。ストライド、パディング、カーネル/フィルターサイズの影響といった概念を理解することも有益です。

AlexNet

AlexNetは深層畳み込みニューラルネットワークで、2012年にAlex Krizhevskyとその同僚によって初めて開発されました。これは、ImageNet LSVRC-2010コンペティションのために画像を分類するように設計されており、その際に最先端の結果を達成しました。このモデルについては、こちらの元の研究論文で詳しく読むことができます。

AlexNetの論文からの重要な要点を振り返りましょう。まず、AlexNetはサイズが(224x224x3)の3チャンネル画像で動作しました。サブサンプリング時には、最大プーリングとReLU活性化を使用しました。畳み込みに使用されるカーネルは11x11、5x5、または3x3のいずれかであり、最大プーリングに使用されるカーネルは3x3のサイズでした。1000クラスに画像を分類しました。また、複数のGPUを利用しました。

データセット

データのロードと前処理から始めましょう。私たちの目的のために、CIFAR-10データセットを使用します。このデータセットは、10クラスの60000枚の32x32カラー画像で構成されており、各クラスには6000枚の画像があります。訓練用の画像は50000枚、テスト用の画像は10000枚です。

ここにデータセット内のクラスと、それぞれのクラスからの10枚のランダムなサンプル画像があります:

ソース: ソース

クラスは完全に相互に排他的です。自動車とトラックの間には重複はありません。「自動車」にはセダン、SUV、その他の類似のものが含まれます。「トラック」には大きなトラックのみが含まれます。どちらにもピックアップトラックは含まれていません。

ライブラリをインポートする

まず、必要なライブラリをインポートし、変数deviceを定義して、ノートブックが利用可能な場合にモデルをGPUでトレーニングすることを知るようにしましょう。

import numpy as np
import torch
import torch.nn as nn
from torchvision import datasets
from torchvision import transforms
from torch.utils.data.sampler import SubsetRandomSampler

Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

データセットをロードしています

torchvision（コンピュータビジョンタスクのためのヘルパーライブラリ）を使用して、データセットを読み込みます。このメソッドには、前処理を非常に簡単かつ直接的にするいくつかのヘルパー関数があります。get_train_valid_loader と get_test_loader の関数を定義し、それを呼び出して CIFAR-10 データを読み込み、処理しましょう:

def get_train_valid_loader(data_dir,
                           batch_size,
                           augment,
                           random_seed,
                           valid_size=0.1,
                           shuffle=True):
    normalize = transforms.Normalize(
        mean=[0.4914, 0.4822, 0.4465],
        std=[0.2023, 0.1994, 0.2010],
    )

    # define transforms
    valid_transform = transforms.Compose([
            transforms.Resize((227,227)),
            transforms.ToTensor(),
            normalize,
    ])
    if augment:
        train_transform = transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ])
    else:
        train_transform = transforms.Compose([
            transforms.Resize((227,227)),
            transforms.ToTensor(),
            normalize,
        ])

    # load the dataset
    train_dataset = datasets.CIFAR10(
        root=data_dir, train=True,
        download=True, transform=train_transform,
    )

    valid_dataset = datasets.CIFAR10(
        root=data_dir, train=True,
        download=True, transform=valid_transform,
    )

    num_train = len(train_dataset)
    indices = list(range(num_train))
    split = int(np.floor(valid_size * num_train))

    if shuffle:
        np.random.seed(random_seed)
        np.random.shuffle(indices)

    train_idx, valid_idx = indices[split:], indices[:split]
    train_sampler = SubsetRandomSampler(train_idx)
    valid_sampler = SubsetRandomSampler(valid_idx)

    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=batch_size, sampler=train_sampler)

    valid_loader = torch.utils.data.DataLoader(
        valid_dataset, batch_size=batch_size, sampler=valid_sampler)

    return (train_loader, valid_loader)


def get_test_loader(data_dir,
                    batch_size,
                    shuffle=True):
    normalize = transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225],
    )

    # define transform
    transform = transforms.Compose([
        transforms.Resize((227,227)),
        transforms.ToTensor(),
        normalize,
    ])

    dataset = datasets.CIFAR10(
        root=data_dir, train=False,
        download=True, transform=transform,
    )

    data_loader = torch.utils.data.DataLoader(
        dataset, batch_size=batch_size, shuffle=shuffle
    )

    return data_loader

CIFAR10 dataset 
train_loader, valid_loader = get_train_valid_loader(data_dir = './data',                                      batch_size = 64,
                       augment = False,                                          random_seed = 1)

test_loader = get_test_loader(data_dir = './data',
                              batch_size = 64)

コードを分解しましょう:

• 私たちは、トレーニング/検証セットとテストセットをそれぞれロードするために、二つの関数get_train_valid_loaderとget_test_loaderを定義します。
• 私たちは、データセット内の各チャネル（赤、緑、青）の平均と標準偏差を用いて変数normalizeを定義することから始めます。これらは手動で計算できますが、CIFAR-10は非常に人気があるため、オンラインでも入手可能です。
• 私たちのトレーニングデータセットには、より堅牢なトレーニングのためにデータセットを拡張するオプションも追加し、画像の数を増やすことができます。注意：拡張はトレーニングサブセットにのみ適用され、評価目的で使用されるバリデーションおよびテストサブセットには適用されません。
• 私たちは、トレーニングデータセットをトレーニングセットと検証セット（90:10の比率）に分割し、全体のトレーニングセットからランダムに部分集合を作成しました。
• バッチサイズを指定し、データセットをロードする際にシャッフルすることで、各バッチに含まれるラベルの種類にばらつきが生まれるようにしています。これにより、最終的なモデルの効力が高まります。
• 最後に、データローダーを使用します。これは、CIFAR-10のような小さなデータセットの場合、パフォーマンスに影響しないかもしれませんが、大きなデータセットの場合にはパフォーマンスを著しく妨げる可能性があり、一般的に良いプラクティスと考えられています。データローダーを使用することで、データをバッチで反復処理でき、データのロードは反復中に行われ、一度にすべてをRAMに取り込むのではありません。

スクラッチからのAlexNet

まずはコードから始めましょう:

class AlexNet(nn.Module):
    def __init__(self, num_classes=10):
        super(AlexNet, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(3, 96, kernel_size=11, stride=4, padding=0),
            nn.BatchNorm2d(96),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size = 3, stride = 2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(96, 256, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size = 3, stride = 2))
        self.layer3 = nn.Sequential(
            nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(384),
            nn.ReLU())
        self.layer4 = nn.Sequential(
            nn.Conv2d(384, 384, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(384),
            nn.ReLU())
        self.layer5 = nn.Sequential(
            nn.Conv2d(384, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size = 3, stride = 2))
        self.fc = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(9216, 4096),
            nn.ReLU())
        self.fc1 = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(4096, 4096),
            nn.ReLU())
        self.fc2= nn.Sequential(
            nn.Linear(4096, num_classes))

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = self.layer5(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        out = self.fc1(out)
        out = self.fc2(out)
        return out

AlexNetモデルの定義

上記のコードがどのように機能するのか見てみましょう:

• PyTorchでCNNかどうかに関わらず、任意のニューラルネットワークを定義する最初のステップは、nn.Moduleを継承するクラスを定義することです。これは、私たちが利用する必要のある多くのメソッドを含んでいます。
• その後、主に2つのステップがあります。まずは、__init__ 内でCNNで使用する層を初期化すること、次にそれらの層が画像を処理する順序を定義することです。これはforward関数内で定義されます。
• アーキテクチャ自体については、まずnn.Conv2D関数を使用して、適切なカーネルサイズと入力/出力チャネルを指定して畳み込み層を定義します。また、nn.MaxPool2D関数を使用して最大プーリングも適用します。PyTorchの良いところは、畳み込み層、活性化関数、および最大プーリングを1つの層に統合できることです（それらは別々に適用されますが、整理に役立ちます）nn.Sequential関数を使用して行います。
• 次に、完全結合層を線形（nn.Linear）、ドロップアウト（nn.Dropout）、およびReLu活性化関数（nn.ReLU）を使用して定義し、これらをnn.Sequential関数と組み合わせます。
• 最後に、私たちの最後の層は10個のニューロンを出力し、これは10種類の物体の最終的な予測となります。

ハイパーパラメータの設定

トレーニングの前に、損失関数やオプティマイザー、バッチサイズ、学習率、エポック数など、いくつかのハイパーパラメータを設定する必要があります。

num_classes = 10
num_epochs = 20
batch_size = 64
learning_rate = 0.005

model = AlexNet(num_classes).to(device)

Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, weight_decay = 0.005, momentum = 0.9)  

Train the model
total_step = len(train_loader)

最初に、シンプルなハイパーパラメータ（エポック、バッチサイズ、学習率）を定義し、クラス数（この場合は10）を引数として使用してモデルを初期化し、モデルを適切なデバイス（CPUまたはGPU）に転送します。それからコスト関数をクロスエントロピー損失として定義し、最適化器をAdamとして定義します。これには多くの選択肢がありますが、これらはモデルと与えられたデータで良好な結果をもたらす傾向があります。最後に、トレーニング中にステップをより良く追跡するために total_step を定義します。

トレーニング

現時点でモデルのトレーニングの準備が整いました:

total_step = len(train_loader)

for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):  
        # Move tensors to the configured device
        images = images.to(device)
        labels = labels.to(device)

        # Forward pass
        outputs = model(images)
        loss = criterion(outputs, labels)

        # Backward and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' 
                   .format(epoch+1, num_epochs, i+1, total_step, loss.item()))

    # Validation
    with torch.no_grad():
        correct = 0
        total = 0
        for images, labels in valid_loader:
            images = images.to(device)
            labels = labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            del images, labels, outputs

        print('Accuracy of the network on the {} validation images: {} %'.format(5000, 100 * correct / total)) 

コードが何をするのか見てみましょう:

• 私たちはエポックの数を繰り返し、その後にトレーニングデータのバッチを繰り返します。
• 私たちは、使用しているデバイス、つまり GPU または CPU に応じて、画像とラベルを変換します。
• フォワードパスでは、モデルを使用して予測を行い、その予測と実際のラベルに基づいて損失を計算します。
• 次に、逆伝搬を行い、実際に重みを更新してモデルを改善します。
• その後、毎回の更新前にoptimizer.zero_grad()関数を使用して勾配をゼロに設定します。
• 次に、loss.backward()関数を使用して新しい勾配を計算します。
• そして最後に、optimizer.step() 関数で重みを更新します。
• また、各エポックの最後に、検証セットを使用してモデルの精度を計算します。この場合、勾配は必要ないため、より速い評価のために with torch.no_grad() を使用します。

出力は次のように見ることができます。

トレーニングロスと検証精度

私たちが見るように、損失は毎エポックごとに減少しており、これは私たちのモデルが本当に学習していることを示しています。この損失は訓練セットのものであり、損失が非常に小さい場合はオーバーフィッティングを示す可能性があります。これが私たちが検証セットも使用している理由です。検証セットでは精度が上昇しているようで、オーバーフィッティングの可能性は低いことを示しています。では、モデルをテストしてその性能を確認しましょう。

テスト

さて、私たちのモデルが見たことのないデータでどのようにパフォーマンスを発揮するかを見てみましょう。

with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
        del images, labels, outputs

    print('Accuracy of the network on the {} test images: {} %'.format(10000, 100 * correct / total))   

これは、私たちの検証目的のためにコードが全く同じであることに注意してください。

モデルを使用し、わずか6エポックの訓練で、検証セットで約78.8%の精度を得られるようです。

精度のテスト

結論

この記事で行ったことをまとめましょう:

• 私たちはAlexNetモデルのアーキテクチャとさまざまな種類のレイヤーを理解することから始めました。
• 次に、torchvisionを使用してCIFAR-10データセットをロードし、前処理を行いました。
• それから、私たちはPyTorchを使って、ゼロからAlexNetモデルを構築しました。
• 最後に、私たちはCIFAR-10データセットでモデルを訓練しテストしました。モデルは最小限の訓練（6エポック）でテストデータセットで良好なパフォーマンスを示しました。

未来の仕事

この記事は、しっかりとした紹介と実践的な経験を提供しますが、さらに探求し、他に何を達成できるかを発見することで、より多くの知識を得られるでしょう。

• 異なるデータセットを使用してみることができます。その一例がCIFAR-100データセットで、これはCIFAR-10データセットの拡張版で、100のクラスがあります。
• 異なるハイパーパラメータを試してみて、モデルにとっての最適な組み合わせを見つけることができます。
• 最後に、データセットからレイヤーを追加または削除して、モデルの能力への影響を確認することができます。