はじめに

前回、tqdmを使って、プログレスバーを表示してみた。

msdd.hatenablog.com

今回は、実際にpytorchを使って、ニューラルネットワークを学習する時の学習の進行度をプログレスバーで表示してみる。

学習時にバーを表示させる

pytorchを使って、ニューラルネットワークを学習させる時に、プログレスバーを表示させて、学習状況を見やすくしてみる。

バージョン

import tqdm
import torch
import torchvision

print(tqdm.__version__)
print(torch.__version__)
print(torchvision.__version__)

出力

4.43.0
1.4.0
0.5.0

tqdmのバージョンを4.38.0？では、学習の途中で無限ループのような感じになった。なので、tqdmのバージョンを上げて、4.43.0で実行すると、きちんと動くようになった。

import文とパラメータ

import文

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms

from tqdm import tqdm

パラメータの変数を宣言する。gpuありで学習した。クラス数は、今回用いたデータセットCIFAR10の10クラス。epoch数は、学習状況を少しみたいだけなので、少なくした。 batch sizeとlrは適当にした。

device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
batch_size=128
num_classes=10
num_epochs=20
lr=0.1

データセットの用意

データセットの用意をする。今回は、画像識別のデータセットのCIFAR10を用いた。データセットを読み込み時に変更するtransformは、画像をtensor形式にするToTensor()のみにした。 CIFAR10()でデータをダウンロードし、使えるようにする。その後、データをまとめて処理できるようにbatch_size分を読み出すdataloaderを作り、データの準備は終わり。

transform={
    "train":transforms.Compose([
        transforms.ToTensor(),
    ]),
    "test":transforms.Compose([
        transforms.ToTensor(),
    ])
}

dataset={
    "train":torchvision.datasets.CIFAR10(root="CIFAR10",train=True,transform=transform["train"],download=True),
    "test":torchvision.datasets.CIFAR10(root="CIFAR10",train=False,transform=transform["test"],download=True)
}

dataloader={
    "train":torch.utils.data.DataLoader(dataset["train"],batch_size=batch_size,shuffle=True,num_workers=2),
    "test":torch.utils.data.DataLoader(dataset["test"],batch_size=batch_size,shuffle=False,num_workers=2)
}

モデル、optimizer、ロス関数を用意

ネットワークはresnet18を使用した。最後の出力の数を、データセットのクラス数の10クラスへと変えて、前に宣言したdeviceへ送る。optimizerはadamでロス関数はcross entropy lossを使用した。

model=torchvision.models.resnet18(pretrained=False)
model.fc=nn.Linear(model.fc.in_features,num_classes)
model=model.to(device)

optimizer=optim.Adam(model.parameters(),lr=lr)
criterion=nn.CrossEntropyLoss()

学習でプログレスバー表示

学習部分のコード。 1epochごとに、学習とテストを行う。データローダーでバッチごとに処理をしていくたびにプログレスバーを更新していくようにしている。

tqdm()でバーを表示できるようにする。今回は手動でバーを更新する方法を使った。引数のtotalに繰り返しの合計数などを入れ、繰り返し時にupdate(n)でバーをnだけ進めて行く。今回はdataloaderのサイズ、バッチごとに処理する回数を入れた。引数のunitは入れた文字で処理速度の単位を変えることができる。今回はバッチごとの処理をしているので、batchを入れた。\ set_description(f"Epoch[{epoch}/{num_epochs}]({phase})")で、バーの前側の部分に現在のepoch数と最終のepoch数、 trainかtestかを表示するようにしている。文字列の前にfを付けることで、{}で囲んだ中に変数を入れることで、変数の値を表示できて便利である。\ set_postfix({"loss":running_loss.item(),"accuracy":accuracy })でバッチごとに学習またはテストした途中の段階でのlossと精度accuracyを更新している。\ 最後に、バッチ1つ分の処理が終了したので、update(1)でバーを1進め、totalになるまで繰り返す。

def train(model,dataloader,otpimizer,criterion,num_epochs,device):

    for epoch in range(1,num_epochs+1):

        for phase in ["train","test"]:

            if phase=="train":
                model.train()
            elif phase=="test":
                model.eval()

            with torch.set_grad_enabled(phase=="train"):
                loss_sum=0
                corrects=0
                total=0

                with tqdm(total=len(dataloader[phase]),unit="batch") as pbar:
                    pbar.set_description(f"Epoch[{epoch}/{num_epochs}]({phase})")
                    for imgs,labels in dataloader[phase]:   
                        imgs,labels=imgs.to(device),labels.to(device)
                        output=model(imgs)
                        loss=criterion(output,labels)

                        if phase=="train":
                            optimizer.zero_grad()
                            loss.backward()
                            optimizer.step()

                        predicted=torch.argmax(output,dim=1) ## dimあってる？
                        corrects+=(predicted==labels).sum()
                        total+=imgs.size(0)
                            
                        #loss関数で通してでてきたlossはCrossEntropyLossのreduction="mean"なので平均
                        #batch sizeをかけることで、batch全体での合計を今までのloss_sumに足し合わせる
                        loss_sum+=loss*imgs.size(0) 

                        accuracy=corrects.item()/total
                        running_loss=loss_sum/total
                        pbar.set_postfix({"loss":running_loss.item(),"accuracy":accuracy })
                        pbar.update(1)

学習を行うと、バーが表示され、lossとaccuracyが更新されているのがわかる。進むにつれてlossが下がっていっており、学習が進んでいるとわかる。予測時間も表示されるので、どれくらいで終わるかも予測しやすい。

train(model,dataloader,optimizer,criterion,num_epochs,device)

f:id:msdd:20200326180450g:plain

まとめ

実際の学習コードでプログレスバーを表示して、学習の進み具合を見れるようにしてみた。実際にlossが下がっているかを確認でき学習の状況を把握しやすくなる。

msdd’s blog

deep learning勉強中。プログラム関連のこと書きます。

PyTorchでの学習の進み具合をプログレスバーで表示してみる