これまで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))
実行するとデータセットの中から無作為に画像データを選び、その画像を出力します。
モデルの読み込み
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:#５回に１回エポックを表示 print('Epoch {}/{}'.format(epoch+1, EPOCH)) print('-' * 10)
モデルの保存
トレーニングが終わったモデルを保存します。 Pytorchで保存したモデルを呼び出すようにしたいときには以下の1行で完了します。
torch.save(net.state_dict(),"detector/face_mask_detector.pth")
失敗談
ここで素人ならでは(?)の失敗談をまとめてみます。
学習結果が発散してる
最初に学習をかけたときにどういうわけか学習結果が波打つように変化したグラフが出力されました。
なんかの信号入力をグラフにしたのかと思うぐらいですw。
これは逆伝搬したときに勾配が更新されることが原因のようです。
これによって、すでに存在している勾配があるとそこに勝手に追加されてしまうため、lossのグラフが発散してしまったようです。
そこでトレーニングの段階でパラメータの更新が終わったらoptimizer.zero_grad()を行うことで勾配情報がリセットされて前のトレーニングの結果に引きずられることがなくなります。
↓こちらの記事が参考になりました
中途半端な学習結果
optimizer.zero_grad()を入れてなんとか理想に近いグラフはできましたが、ノイズも混ざって中途半端なグラフになりました。
Validに関しては1に到達しているものもあります(1に到達せず収束するのが理想)。
これはデータ数を増やせば解決します。
各クラス1600枚ほど用意するとこんな感じで激しいノイズを起こさずにきれいに収束したグラフをできました。
ソースコード
今回作ったソースコードはこちらのGithubレポジトリで公開しています。
まとめ
今回は今までチュートリアル止まりで終わっていたディープラーニングを自分で実装してみました。 難しそうに見えたものもライブラリの力で簡単に実装できるのはいいなと思いました。 特にPytorchは導入もシンプルでドキュメントも割と丁寧なのでこれからどんどん使いこなしたいですね。
