Pytorchでcifar10画像分類。Colabで

前回はpytorchチュートリアルを参考にcifar10でのクラス分類をUbuntu18のノートPC上で実施してみましたが、実際に使える学習データを作るにはやっぱGPUが必要なわけですが、アマチュアには予算、スペース的に厳しいわけです。。

そこで最近話題のGoogle様がリリースした新サービスColab上にて私もGPUで学習したいと思い試してみました。これ前回の記事です。

Colabを使う上でよく理解しとかないといけないのが「12時間ルール」ですよね(2020年3月での状況です)。仮想マシンがリセットされると。マシンがリセットされる、というのは実際使ってみると分かるのですが、完全にOSが初期化されちゃう、ってことですよね。例えばpipコマンドでインストールしたパッケージなんかも全部無くなっちゃうと。自分で追加したコードも全部無くなっちゃいます。ちなみにpytorchはデフォルトでインストールされてますので、バージョンにこだわりがなければそのまま使えますね。

pipでインストールしたものが消えるならまた入れればよいのですが、自分で作ったものが消えちゃったりしたら悲劇ですので、気を付けてください。。

ですので現状使い方としては、仮想マシンからGoogleDriveをマウント出来ますので、ある程度pythonファイルをパッケージ対応をしておいてGoogleDriveに置いて、実行するColabファイルにてsys.path.appendでパスを付けといて、仮想マシン上で実行する、という流れになりますね。Colabファイル自体もGoogleDrive上に保存しますが、実行すると仮想マシン上に自動でロードされる感じでしょうか。

早速やってみます。

前回同様ファイルは3枚ですが(model.py,train.py,test.py)、パッケージに対応&GoogleDrive&仮想PCで絶対パスを渡したいので、ちょっと手を加えています。あとGPU対応ですね。冗長ですが上げときます。

GoogleColabに上げる前にパッケージとして動作するか、ローカル環境でJupyterNotebookで試すと良いですが、ここでは省略です。

#model.py
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
#train.py
import torch
import torchvision
import torchvision.transforms as transforms

import torch.optim as optim
import torch.nn as nn

from . import model


def run(hozonsaki):
    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    net = model.Net().to(device)

    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                            download=True, transform=transform)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                              shuffle=True, num_workers=2)

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

    for epoch in range(2):  # loop over the dataset multiple times

        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            # get the inputs; data is a list of [inputs, labels]
            inputs, labels = data
            inputs=inputs.to(device)
            labels=labels.to(device)

            # zero the parameter gradients
            optimizer.zero_grad()

            # forward + backward + optimize
            outputs = net(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            # print statistics
            running_loss += loss.item()
            if i % 2000 == 1999:    # print every 2000 mini-batches
                print('[%d, %5d] loss: %.3f' %
                      (epoch + 1, i + 1, running_loss / 2000))
                running_loss = 0.0

    print('Finished Training')

    PATH = hozonsaki+'/cifar_net.pth'
    torch.save(net.state_dict(), PATH)
#test.py
import torch
import torchvision
import torchvision.transforms as transforms

from . import model
import matplotlib.pyplot as plt
import numpy as np


def imshow(img):
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


def run(hozonsaki):
    PATH = hozonsaki+'/cifar_net.pth'
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    net = model.Net().to(device)
    net.load_state_dict(torch.load(PATH))

    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])


    testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                           download=True, transform=transform)
    testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                             shuffle=False, num_workers=2)

    classes = ('plane', 'car', 'bird', 'cat',
               'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

    dataiter = iter(testloader)
    images, labels = dataiter.next()
    

    # print images
    imshow(torchvision.utils.make_grid(images))
    print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(4)))

    images=images.to(device)
    labels=labels.to(device)

    outputs = net(images)
    _, predicted = torch.max(outputs, 1)

    print('Predicted: ', ' '.join('%5s' % classes[predicted[j]]
                                  for j in range(4)))


前回とほとんど変更はないのですが、大事そうな変更点です。。

import model

から

from . import model

です。implicit relative import からexplicit relative importに変更してあります。パッケージ内ではimplicit relative import が使えません。

あとはcuda利用する点ですね。

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

でモデルとデータにcuda対応ですね。

net = model.Net().to(device)
inputs=inputs.to(device)
labels=labels.to(device)

GoogleDriveにアップする

以上の3ファイルをGoogleDriveにアップしてみます。以下のような構成でアップしました。

GoogleDriveにパッケージアップロード

実際は実行後なのでキャッシュが残ってますがご容赦ください。

でこれらを呼び出すColabファイル(JupyterNotebookファイル)を作成します。

train.pyを呼び出すファイルCallTestCifarTest.ipynbとtest.pyを呼び出すCallTestCifarTest.ipynbです。まあテストはあまりColabでやる意味はないですが、一応やってみます。

GoogleDriveにColabファイル作成

Colabで学習部分を呼び出してみる。

.ipynbファイルを開くと左側に仮想マシンのディレクトが見えますね。ここに「Mount Drive」ボタンがありますので、これをマウントすると仮想マシンからGoogleDriveが見えます。

ColabでGoogleDriveマウント

このように先程作ったファイルがみえますね。

ColabでGoogleDriveマウントした後

.ipynbファイルを作る前にGPUを有効にしといたほうがよいですね。「Runtime」ー>「Change runtime type」でCPU(None),GPU,TPUが選べますね。ここではGPUです。それぞれの.ipynbファイルで設定してやらないとダメみたいですね。

GPU設定

JupyterNotebookファイル

パッケージへの絶対パスはここでは’/content/drive/My Drive/CifarTest’になります。
以下呼び出し用の.ipynbファイルです。

#CallTrainCifarTest.ipynb
import sys
sys.path.append('/content/drive/My Drive/CifarTest')

from CifarTest.train import run

run('/content/drive/My Drive/CifarTest')

実行するとこんなですね。

学習呼び出しFromColab

無事終わるとGoogleDriveに学習済みファイルが出来ているかと思います。

/content/drive/My Drive/CifarTest/cifar_net.pth

推論も同様に動くか試してみます。

#CallTestCifarTest.ipynb
import sys
sys.path.append('/content/drive/My Drive/CifarTest')

from CifarTest.test import run

run('/content/drive/My Drive/CifarTest')
推論呼び出しFromColab

となんとか動作するのではと思いますが、どうでしょうか。

GoogleDriveー>仮想マシンへの更新は即時ではない

GoogleDrive内ファイルのパッケージを変更しても、仮想マシンへの反映がすぐではありませんね。間違いないのは「Restart Runtime」や「Factory Reset Runtime」を試せば更新されるようです。