PyTorchでマスクしている人を分類してみた

これまでPythonをメインの言語としてプライベートで開発をしていましたが、だいたいWebアプリやIoT系の開発がメインになってしまい、機械学習やDeepLearningは最初のチュートリアルをちょっと触った程度で終わっていました。 Pythonやっててそれはさすがにまずいと思い、自分でDeeplearningで何か実装してみようと思っていました。 そこで今回はマスクを付けている人をPyTorchで分類するモデルを作ってみようと思います。

データセット

まず肝心なデータセットの用意をします。 それぞれ以下のサイトから入手しました。

cabani/MaskedFace-Net - マスクをつけた顔画像のデータセットを取得

Flickr-Faces-HQ Dataset (FFHQ) - マスクをつけていない顔画像のデータセットを取得(↑のデータセットはこのデータセットをベースにマスクをつけているっぽい)

Datasetクラスを作る

PyTorchでは自前のデータセットを扱うにはDatasetクラスを用意する必要がありますが、用意したデータセットを以下のディレクトリ構成で保存すれば、PyTorchのImageFolderを使うことができます。 dataディレクトリを参照することで各クラスのデータセットを取り込みます。

data
├─without_mask
└─with_mask

データセットTensor型に変換するtransformを合わせると以下のようにDatasetクラスを定義できます。

from torchvision import transforms
from torchvision.datasets import ImageFolder
# データセットの読み込み
data_transforms = {
    'train': transforms.Compose([
        transforms.Resize(size=(224,224)),
        transforms.RandomRotation(degrees=15),
        transforms.ColorJitter(),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Resize(size=(224,224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

dataset = ImageFolder("data", data_transforms['train'])
dataset.class_to_idx

dataset.class_to_idxでクラス名、IDの一覧を辞書で確認できます。 以下のようにクラス名、IDが割り振られていたらOKです。

{'with_mask': 0, 'without_mask': 1}

読み込んだデータを訓練用、評価用でそれぞれ分けます。 (訓練用8割、評価用2割)

# 訓練用、評価用にデータを分ける
all_size = len(dataset)
train_size = int(0.8 * all_size)
val_size = all_size - train_size
dataset_size  = {"train":train_size, "val":val_size}
train_data, val_data = random_split(dataset, [train_size, val_size])

DataLoaderの定義とデータの中身の確認

次にDataLoaderを定義します。

# DataLoaderを定義
batch_size = 60
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_data, batch_size=batch_size, shuffle=False)
dataloaders = {'train':train_loader, 'val':val_loader}

ここで読み込んだデータの中身を確認します。 ただデータ量が多いので、torchvision.utils.make_grid(images, nrow=10)を使って表示数を絞り込みます。 データの中身を見るときにはmatplotlibをつかいます。

dataiter = iter(dataloaders['train'])
images, labels = dataiter.next()

# Make a grid from batch
out = torchvision.utils.make_grid(images, nrow=10)
plt.imshow(out.permute(1, 2, 0))

実行するとデータセットの中から無作為に画像データを選び、その画像を出力します。

f:id:KMiura:20210118024825p:plain

モデルの読み込み

torch.hub.loadを使うとトレーニング済みのモデルを読み込むことができます。 今回はこれを使ってMobileNetV2を読み込みたいと思います。

Model = torch.hub.load('pytorch/vision:v0.6.0', 'mobilenet_v2', pretrained=True)

転移学習

このモデルはすでに既存のデータで学習しているため、今回用意したデータを学習させるためには転移学習を行います。 以下のようにパラメータを取り出して、各層をフリーズさせます。

for param in Model.parameters():
    param.requires_grad = False

ここで今回用意したデータを学習させるために、最終層のネットワークを上書きします。

# 今回の学習用に最終層のネットワークを上書き
Model.classifier[1] = nn.Sequential(
                      nn.Linear(1280, 256),
                      nn.ReLU(),
                      nn.Linear(256, 128),
                      nn.ReLU(), 
                      nn.Dropout(0.4),
                      nn.Linear(128, 64),
                      nn.ReLU(),
                      nn.Linear(64, 32),
                      nn.ReLU(),
                      nn.Dropout(0.4),
                      nn.Linear(32, 2), 
                      nn.LogSoftmax(dim=1))

ロス関数とoptimizerを定義

cross-emtropyとSGDを定義します。

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

レーニン

用意したネットワークとotimizerを使ってトレーニングを行います。

for epoch in range(EPOCH):
    for phase in ['train', 'val']:
        if phase == 'train':
            model.train()
        else:
            model.eval()

        running_loss = 0.0
        running_corrects = 0

        for inputs, labels in dataloaders[phase]:
            inputs = inputs.to(device)
            labels = labels.to(device)
            
            with torch.set_grad_enabled(phase == 'train'):
                outputs = model(inputs)
                _, preds = torch.max(outputs, 1)
                loss = criterion(outputs, labels)

                # backward + optimize only if in training phase
                if phase == 'train':
                    optimizer.zero_grad()
                    loss.backward()
                    optimizer.step()
            running_loss += loss.item() * inputs.size(0)
            correct = torch.eq(torch.max(F.softmax(outputs, dim=1), dim=1)[1], labels).view(-1)
            running_corrects += torch.sum(correct).item()
        epoch_loss = running_loss / dataset_size[phase]
        epoch_acc = running_corrects / dataset_size[phase]
        loss_dict[phase].append(epoch_loss)
        acc_dict[phase].append(epoch_acc)
        
        print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
        
        if phase == 'val' and epoch_loss<0.005:
            best_loss = epoch_loss
            best_acc = epoch_acc
            best_model_wts = copy.deepcopy(model.state_dict())

    if (epoch+1)%5 == 0:#5回に1回エポックを表示
        print('Epoch {}/{}'.format(epoch+1, EPOCH))
        print('-' * 10)

モデルの保存

レーニングが終わったモデルを保存します。 Pytorchで保存したモデルを呼び出すようにしたいときには以下の1行で完了します。

torch.save(net.state_dict(),"detector/face_mask_detector.pth")

失敗談

ここで素人ならでは(?)の失敗談をまとめてみます。

学習結果が発散してる

f:id:KMiura:20210119032520p:plain

最初に学習をかけたときにどういうわけか学習結果が波打つように変化したグラフが出力されました。 なんかの信号入力をグラフにしたのかと思うぐらいですw。 これは逆伝搬したときに勾配が更新されることが原因のようです。 これによって、すでに存在している勾配があるとそこに勝手に追加されてしまうため、lossのグラフが発散してしまったようです。 そこでトレーニングの段階でパラメータの更新が終わったらoptimizer.zero_grad()を行うことで勾配情報がリセットされて前のトレーニングの結果に引きずられることがなくなります。

↓こちらの記事が参考になりました

ohke.hateblo.jp

中途半端な学習結果

f:id:KMiura:20210119034024p:plain optimizer.zero_grad()を入れてなんとか理想に近いグラフはできましたが、ノイズも混ざって中途半端なグラフになりました。 Validに関しては1に到達しているものもあります(1に到達せず収束するのが理想)。 これはデータ数を増やせば解決します。 各クラス1600枚ほど用意するとこんな感じで激しいノイズを起こさずにきれいに収束したグラフをできました。 f:id:KMiura:20210119034939p:plain

ソースコード

今回作ったソースコードはこちらのGithubレポジトリで公開しています。

github.com

まとめ

今回は今までチュートリアル止まりで終わっていたディープラーニングを自分で実装してみました。 難しそうに見えたものもライブラリの力で簡単に実装できるのはいいなと思いました。 特にPytorchは導入もシンプルでドキュメントも割と丁寧なのでこれからどんどん使いこなしたいですね。