Python3
PyTorch

PyTorch入門 メモ

この記事に関して

  • PyTorch TutorialのGETTING STARTEDで気になったところのまとめ
  • 数学的な話は省略気味

PyTorchによる深層学習

PyTorchとは

  • 柔軟性と速度を兼ね備えた深層学習のプラットフォーム
  • GPUを用いた高速計算が可能なNumpyのndarraysと似た行列表現tensorを利用可能
  • pipを用いる場合は以下でインストール
pip install torch torchvision
  • 以下でimport
import torch

tensorの初期化&プロパティ

x = torch.empty(5, 3, dtype=torch)         # 0埋め
x = x.new_ones(3, 2, dtype=torch.double)   # 既存のtensorの型変換&1埋め
x = torch.rand(4, 2)                       # 乱数
x = torch.randn_like(x, dtype=torch.float) # 既存のtensorを乱数で埋める
x = torch.tensor([5.5, 3])                 # データから直接構成
x.size()                                   # サイズ取得

tensorの基本操作

  • 加算
    • 他の演算も大体同じ
    • 以下のzresultは等価だが,resultの形式の場合は事前に初期化する必要あり
    • _で終わるメソッドは呼び出し元の変数の値を変化させる
      • x.copy_(y), x.t_() など
z = x+y                     # xとyはtensorで同じ型
result = torch.empty(4, 2)  # xとyと同じ型
torch.add(x, y, out=result) # result == z
y.add_(x)                   # yにxを加算
  • Indexing

例えば,xが以下のtensorであるとき,

tensor([[0.8765, 0.9070],
        [0.8698, 0.6337],
        [0.7680, 0.8755],
        [0.6074, 0.3790]])

以下のようにIndexing可能

x[:, 0] # tensor([0.8765, 0.8698, 0.7680, 0.6074])
x[2, :] # tensor([0.7680, 0.8755])
  • Resize
    • -1で埋めた箇所は他の値から勝手に推測してくれる
x = torch.randn(4, 4) # torch.Size([4, 4])
y = x.view(16)        # torch.Size([16])
z = x.view(-1, 8)     # torch.Size([2, 8])
  • 要素が1個のtensorから値の取り出し
x.item()
  • その他の操作はここ読んで

tensorとNumPy配列の変換

ポインタを渡しているような挙動をする

a = torch.ones(5)   # tensor([1., 1., 1., 1., 1.])
b = a.numpy()       # array([1., 1., 1., 1., 1.], dtype=float32)

a.add_(1)
print(a)            # tensor([2., 2., 2., 2., 2.])
print(b)            # [2. 2. 2. 2. 2.]

np.add(b, 1, out=b)
print(a)            # tensor([3., 3., 3., 3., 3.])
print(b)            # [3. 3. 3. 3. 3.]

CUDA Tensors

  • GPUで動くよ! (CUDA Support)
    • to()を用いて明示的にCPU-GPU間を移動する必要があるらしい
if torch.cuda.is_available():
    device = torch.device("cuda")
    y = torch.ones_like(x, device=device)
    x = x.to(device)
    z = x+y
    print(z)                          # tensor([0.8696], device='cuda:0')
    print(z.to("cpu", torch.double))  # tensor([0.8696], dtype=torch.float64)

Neural Networks

ネットワークの定義

  • LeNetの例
class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        # kernel
        self.conv1 = nn.Conv2d(1,  6, 5) # 1 input channel, 6 output channels, 5x5 square convolution
        self.conv2 = nn.Conv2d(6, 16, 5)

        # affine operation
        self.fc1 = nn.Linear(16*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # Max pooling over a (2, 2) window
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2) # sizeが同じなら片方省いていいらしい
        x = x.view(-1, self.num_flat_features(x))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]
        num_features = 1
        for s in size:
            num_features *= s
        return num_features

net = Net()

ランダムな画像の順伝播

input = torch.randn(1, 1, 32, 32)
out = net(input)
print(out)   # tensor([[-0.0285, -0.0103,  0.1103,  0.0361, -0.0667, -0.0253, -0.0585,  0.0538, 0.0326, -0.1365]], grad_fn=<ThAddmmBackward>)

勾配初期化&勾配計算

net.zero_grad()                  # ネットワークの勾配をすべて0で初期化
out.backward(torch.randn(1, 10)) # 乱数を入力として勾配計算

Loss関数

output = net(input)              # 入力データをネットワークに突っ込む
target = torch.randn(10)         # テキトーな値を正解とする
target = target.view(1, -1)      # 整形
criterion = nn.MSELoss()         # ロス関数定義
loss = criterion(output, target) # ロス計算

Lossを用いたBackprop計算

net.zero_grad()
print(net.conv1.bias.grad) # tensor([0., 0., 0., 0., 0., 0.])
loss.backward()
print(net.conv1.bias.grad) # tensor([0.0020, 0.0184, 0.0053, 0.0086, 0.0065, 0.0103])
  • その他のModuleやらLoss関数やら知りたいときは,ここ読んで

重みの更新

  • シンプルな実装
learning_rate = 0.01
for f in net.parameters():
    f.data.sub_(f.grad.data * learning_rate)
  • 最適化ver.
import torch.optim as optim

# create your optimizer
optimizer = optim.SGD(net.parameters(), lr=0.01)

# 以下をデータ毎にループする
optimizer.zero_grad()   # zero the gradient buffers
output = net(input)
loss = criterion(output, target)
loss.backward()
optimizer.step()    # Does the update

Training a classifier

  • PyTorchを用いて何らかの学習をしたいときは,事前にNumPy配列として読み込んでからtensorに変換するのが一般的
    • 画像: torchvision, PillowやOpenCV
    • 音響: scipyやlibrosa
    • 文章: デフォルトpython, Cython, NLTK, SpaCy
  • ここではtorchvisionを用いたCIFAR10の学習について触れる

1. データのロードと正規化

# tensorへの変換と正規化器
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

# trainingデータのロード
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)

# testデータのロード
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")
画像とラベルの表示
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)))

# get some random training images
dataiter = iter(trainloader)
images, labels = dataiter.next()

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

2. CNNの定義

  • ↑のLeNetとほぼ同じなので説明省略
  • 入力が3チャンネルになっているところだけ注意
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) # ここが3チャンネル
        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
net = Net()

3. Loss関数とOptimizerの定義

CrossEntropyとMomentum SGDを使用

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

4. ネットワークの学習

for epoch in range(2):  # すべてのデータを2回学習(epochs=2)

    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # get the inputs
        inputs, labels = data

        # 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

5.テストデータを用いて評価

画像と推定結果と正解ラベルの表示
dataiter = iter(testloader)
images, labels = dataiter.next()
outputs = net(images)
_, predicted = torch.max(outputs, 1)
print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(4)))
print('Predicted: ', ' '.join('%5s' % classes[predicted[j]]
                              for j in range(4)))
imshow(torchvision.utils.make_grid(images))
画像全体でのAccuracyの測定
  • with torch.no_grad():によりパラメータの保存をストップしている
correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print("Accuracy of the network on the 10000 test images: %d %%" % (100*correct/total))
クラス毎のAccuracyの測定
class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predicted = torch.max(outputs, 1)
        c = (predicted == labels).squeeze()
        for i in range(4):
            label = labels[i]
            class_correct[label] += c[i].item()
            class_total[label] += 1


for i in range(10):
    print('Accuracy of %5s : %2d %%' % (
        classes[i], 100 * class_correct[i] / class_total[i]))

GPUを用いた並列処理

  • 以下のようにモデルを簡単にGPUに移動できる
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
  • tensorも同様にGPUに移動できる
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
gpu_tensor = cpu_tensor.to(device)
  • デフォルトではGPUは1個しか使うことができないため,DataParallelを用いてモデルを並列に動かす必要がある
model = nn.DataParallel(model)

GPUを複数用いた学習の例

ImportとParametersの定義
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

input_size = 5
output_size = 2
batch_size = 30
data_size = 100

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
ダミーデータセットの定義
  • 自前のデータセットを表現するためにはabstract classであるtorch.utils.data.Datasetを継承し,次の2つのメソッドをoverrideする必要がある
    • __len__: データセットのサイズ
    • __getitem__: dataset[i]のようなindexingを行うために定義が必要
class RandomDataset(Dataset):

    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len

rand_loader = DataLoader(dataset=RandomDataset(input_size, data_size),
                         batch_size=batch_size, shuffle=True)
Simpleなモデルの定義
class Model(nn.Module):
    def __init__(self, input_size, output_size):
        super(Model, self).__init__()
        self.fc = nn.Linear(input_size, output_size)

    def forward(self, input):
        output = self.fc(input)
        print("\tIn Model: input size", input.size(), "output size", output.size())
        return output
モデルの作成とDataParallelを用いた複数GPUの使用宣言
model = Model(input_size, output_size)
if torch.cuda.device_count() > 1:
    print(torch.cuda.device_count(), "GPUs")
    model = nn.DataParallel(model)
model.to(device)
モデルの学習
for data in rand_loader:
    input = data.to(device)
    output = model(input)
    print("input size", input.size(), "output_size", output.size())

データのロードとAugmentation

Import

import os
import torch
import pandas as pd
from skimage import io, transform
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms,utils
import warnings
warnings.filterwarnings("ignore")
plt.ion()

データのロード

  • ここから顔のlandmarksのデータセットを取得してテキトーなところに置く
    • このデータセットは画像名image_nameとlandmarks 計68点のx座標及びy座標part_n_x, part_n_yで構成される
landmarks_frame = pd.read_csv("./data/faces/face_landmarks.csv")
n = 65
img_name = landmarks_frame.iloc[n, 0]
landmarks = landmarks_frame.iloc[n, 1:].as_matrix()
landmarks = landmarks.astype("float").reshape(-1, 2)

Dataset Classの定義

自前のデータセットを使うには,↑に書いたように__len____getitem__の定義が必要

class FaceLandmarksDataset(Dataset):

    def __init__(self, csv_file, root_dir, transform=None):
        self.landmarks_frame = pd.read_csv(csv_file)
        self.root_dir = root_dir
        self.transform = transform

    def __len__(self):
        return len(self.landmarks_frame)

    def __getitem__(self, idx):
        img_name = os.path.join(self.root_dir,
                                self.landmarks_frame.iloc[idx, 0])
        image = io.imread(img_name)
        landmarks = self.landmarks_frame.iloc[idx, 1:].as_matrix()
        landmarks = landmarks.astype('float').reshape(-1, 2)
        sample = {'image': image, 'landmarks': landmarks}

        if self.transform:
            sample = self.transform(sample)

        return sample

Transformクラスの定義

  • データセット中の画像が同じサイズでない場合にはサイズを揃える必要がある
  • 以下の3つの変換を駆使して画像変形を行う
    • Rescale: 画像のスケール変換
    • RandomCrop: Data Augmentationのために画像をランダムでクロップする
    • ToTensor: NumPy画像をTorch画像に変換する
      • Numpy画像とTorch画像は表現形式が異なるので注意
  • ↑3つのTransformを定義する際には,関数ではなくクラスを用いるとパラメータの無駄な初期化を減らすことができるため良い
class Rescale(object):
    """Rescale the image in a sample to a given size.

    Args:
        output_size (tuple or int): Desired output size. If tuple, output is
            matched to output_size. If int, smaller of image edges is matched
            to output_size keeping aspect ratio the same.
    """
    def __init__(self, output_size):
        assert isinstance(output_size, (int, tuple))
        self.output_size = output_size

    def __call__(self, sample):
        image, landmarks = sample['image'], sample['landmarks']
        h, w = image.shape[:2]
        if isinstance(self.output_size, int):
            if h > w:
                new_h, new_w = self.output_size * h / w, self.output_size
            else:
                new_h, new_w = self.output_size, self.output_size * w / h
        else:
            new_h, new_w = self.output_size
        new_h, new_w = int(new_h), int(new_w)
        img = transform.resize(image, (new_h, new_w))

        # h and w are swapped for landmarks because for images,
        # x and y axes are axis 1 and 0 respectively
        landmarks = landmarks * [new_w / w, new_h / h]
        return {'image': img, 'landmarks': landmarks}
class RandomCrop(object):
    """Crop randomly the image in a sample.

    Args:
        output_size (tuple or int): Desired output size. If int, square crop
            is made.
    """
    def __init__(self, output_size):
        assert isinstance(output_size, (int, tuple))
        if isinstance(output_size, int):
            self.output_size = (output_size, output_size)
        else:
            assert len(output_size) == 2
            self.output_size = output_size

    def __call__(self, sample):
        image, landmarks = sample['image'], sample['landmarks']
        h, w = image.shape[:2]
        new_h, new_w = self.output_size
        top = np.random.randint(0, h - new_h)
        left = np.random.randint(0, w - new_w)

        image = image[top: top + new_h,
                      left: left + new_w]

        landmarks = landmarks - [left, top]
        return {'image': image, 'landmarks': landmarks}
class ToTensor(object):
    """Convert ndarrays in sample to Tensors."""
    def __call__(self, sample):
        image, landmarks = sample['image'], sample['landmarks']

        # swap color axis because
        # numpy image: H x W x C
        # torch image: C X H X W
        image = image.transpose((2, 0, 1))
        return {'image': torch.from_numpy(image),
                'landmarks': torch.from_numpy(landmarks)}

Transformの適用

  • 複数のTransformクラスは`torchvision.transforms.Composeで合成できる
scale = Rescale(256)
crop = RandomCrop(128)
composed = transforms.Compose([Rescale(256), RandomCrop(224)])

# Apply each of the above transforms on sample.
fig = plt.figure()
sample = face_dataset[65]
for i, tsfrm in enumerate([scale, crop, composed]):
    transformed_sample = tsfrm(sample)

    ax = plt.subplot(1, 3, i + 1)
    plt.tight_layout()
    ax.set_title(type(tsfrm).__name__)
    show_landmarks(**transformed_sample)

plt.show()

データセットのロード(Augmentation込)

  • torch.utils.data.Datasetの引数transformを用いると簡単にデータセットを構成できる
transformed_dataset = FaceLandmarksDataset(csv_file="./data/faces/face_landmarks.csv",
                                           root_dir="./data/faces/",
                                           transform=transforms.Compose([
                                               Rescale(256),
                                               RandomCrop(224),
                                               ToTensor()]))
  • なお,実は基本的なtransformはPyTorchの内部で定義されている(次章)

データのIteration

  • torch.utils.data.DataLoaderを用いると,入力のすべての特徴に対するループを簡単に書ける
dataloader = DataLoader(transformed_dataset, batch_size=4,
                        shuffle=True, num_workers=4)

for i_batch, sample_batched in enumerate(dataloader):
    print(i_batch, sample_batched['image'].size(),
          sample_batched['landmarks'].size())
    ...

転移学習

  • データセット足りてない時には転移学習しようぜ
  • 転移学習には大きく2つのやり方がある
    • ネットワークを乱数で初期化する代わりに事前に学習したネットワークで初期化する(Finetuning)
    • 最終層以外の重みの更新を停止(Freeze)し,最終層だけはテキトーな乱数で初期化する

データのロード

  • ここからデータをロードしてテキトーなところに置く
  • アリとハチを区別するモデルの学習を行う
data_transforms = {
    "train": transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.215])
    ]),
    "val": transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

data_dir = "./data/hymenoptera_data"
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in ["train", "val"]}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4, shuffle=True, num_workers=4) for x in ["train", "val"]}
dataset_sizes = {x: len(image_datasets[x]) for x in ["train", "val"]}
class_names = image_datasets["train"].classes

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

Augmentation結果の確認

def imshow(inp, title=None):
    inp = inp.numpy().transpose((1, 2, 0))
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    inp = std * inp + mean
    inp = np.clip(inp, 0, 1)
    plt.imshow(inp)
    if title is not None:
        plt.title(title)
    plt.pause(0.001)  # pause a bit so that plots are updated


# Get a batch of training data
inputs, classes = next(iter(dataloaders['train']))

# Make a grid from batch
out = torchvision.utils.make_grid(inputs)

imshow(out, title=[class_names[x] for x in classes])

モデルの訓練

  • 以下はモデルを訓練するための関数
  • schedulertorch.optim.lr_schedulerで学習率の調整を行うオブジェクト
def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
    since = time.time()

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    for epoch in range(num_epochs):
        print("Epoch {}/{}".format(epoch, num_epochs-1))
        print("-"*10)

        # Each epoch has a training and validation phase
        for phase in ["train", "val"]:
            if phase == "train":
                scheduler.step()
                model.train()
            else:
                model.eval()

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                optimizer.zero_grad()

                # forward
                with torch.set_grad_enabled(phase=="train"):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # backward + optimize
                    if phase=="train":
                        loss.backward()
                        optimizer.step()

                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds==labels.data)

            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]

            print("{} Loss: {:.4f} Acc: {:.4f}".format(phase, epoch_loss, epoch_acc))

            # deep copy the model
            if phase=="val" and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())

        print()

    time_elapsed = time.time()-since
    print("Training complete in {:.0f}m {:.0f}s".format(time_elapsed // 60, time_elapsed%60))
    print("Best val Acc: {:4f}".format(best_acc))

    # load best model weights
    model.load_state_dict(best_model_wts)
    return model

学習結果の可視化関数の定義

def visualize_model(model, num_images=6):
    was_training = model.training
    model.eval()
    images_so_far = 0
    fig = plt.figure()

    with torch.no_grad():
        for i, (inputs, labels) in enumerate(dataloaders["val"]):
            inputs = inputs.to(device)
            labels = labels.to(device)

            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)

            for j in range(inputs.size()[0]):
                images_so_far += 1
                ax = plt.subplot(num_images//2, 2, images_so_far)
                ax.axis("off")
                ax.set_title("predicted: {}".format(class_names[preds[j]]))
                imshow(inputs.cpu().data[j])

                if images_so_far == num_images:
                    model.train(mode=was_training)
                    return

        model.train(mode=was_training)

Finetuning

学習済ネットワークのロード

model_ft = models.resnet18(pretrained=True)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, 2) 

if torch.cuda.device_count() > 1:
    model_ft = nn.DataParallel(model_ft)

model_ft = model_ft.to(device)
criterion = nn.CrossEntropyLoss()
optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)

# Decay LR by a factor of 0.1 every 7 epochs
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)

Train and evaluate

model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=25)

学習結果の可視化

visualize_model(model_ft)

最終層以外の重みをFreezeする場合

  • 重みの更新を停止するためには,requires_gradFalseにする必要がある(詳細)

学習済ネットワークのロード

model_conv = models.resnet18(pretrained=True)
for param in model_conv.parameters():
    param.requires_grad = False

num_ftrs = model_conv.fc.in_features
model_conv.fc = nn.Linear(num_ftrs, 2)

if torch.cuda.device_count() > 1:
    model_conv = nn.DataParallel(model_conv)

model_conv = model_conv.to(device)

criterion = nn.CrossEntropyLoss()
optimizer_conv = optim.SGD(model_conv.parameters(), lr=0.001, momentum=0.9)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_conv, step_size=7, gamma=0.1)

Train and evaluate

model_conv = train_model(model_conv, criterion, optimizer_conv, exp_lr_scheduler, num_epochs=25)

学習結果の可視化

visualize_model(model_conv)

plt.ioff()
plt.show()

--

モデルのセーブとロード

セーブとロードでは主に以下を使う

  • torch.save: pickleを用いて,serialized objects (=models, tensors, dictionaries)を保存する
  • torch.load: pickleを用いてserialized objectsをロードする
  • torch.nn.Module.load_state_dict: モデルのパラメタが保存されたディクショナリ(state_dict, 後述)をロードする

state_dictとは何か?

以下のようなmodeloptimizerを考える

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

class TheModelClass(nn.Module):
    def __init__(self):
        super(TheModelClass, 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

model = TheModelClass()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

このとき,modeloptimizerstate_dictは以下のようになる

print("Model's state_dict:")
for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())
Model's state_dict:
conv1.weight     torch.Size([6, 3, 5, 5])
conv1.bias   torch.Size([6])
conv2.weight     torch.Size([16, 6, 5, 5])
conv2.bias   torch.Size([16])
fc1.weight   torch.Size([120, 400])
fc1.bias     torch.Size([120])
fc2.weight   torch.Size([84, 120])
fc2.bias     torch.Size([84])
fc3.weight   torch.Size([10, 84])
fc3.bias     torch.Size([10])
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
    print(var_name, "\t", optimizer.state_dict()[var_name])
Optimizer's state_dict:
state    {}
param_groups     [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [140123409333992, 140123409330680, 140123409334064, 140123409330896, 140123409331688, 140123409331832, 140123409331904, 140123409331760, 140123409331616, 140123409331544]}]

推論のためのモデルのセーブとロード

(推奨) state_dictのセーブとロード

  • モデル全体を保存するよりは,パラメタだけ保存した方が効率的
  • PyTorchでモデルを保存するときは,.pt.pthで保存するのが一般的
セーブ
torch.save(model.state_dict(), PATH)
ロード
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()

モデル全体のセーブとロード

  • ロードしたモデルのデータが特定のクラスでしか使えない欠点あり
  • 推論を行う前に,DropoutレイヤとBNレイヤをevaluationモードにするために,model.evalを実行する必要がある
セーブ
torch.save(model, PATH)
ロード
model = torch.load(PATH)
model.eval()

推論もしくは学習を再開するためのチェックポイントのセーブとロード

  • PyTorchではチェックポイントは.tarでまとめるのが一般的
  • ロードする前に初期化しておくこと
  • やはり推論を行う前にはmodel.eval()を行う必要あり
セーブ
torch.save({
          'epoch': epoch,
          'model_state_dict': model.state_dict(),
          'optimizer_state_dict': optimizer.state_dict(),
          'loss': loss,
          ...
          }, PATH)
ロード
model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)

checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

model.eval()
# - or -
model.train()

複数のモデルを一つのファイルにセーブ

  • GANやSeq2Seq,モデルのアンサンブルのように複数のモデルをセーブするときには,ModelとOptimizerのstate_dictをそれぞれ同一ファイルに保存する
  • PyTorchでは.tarでまとめるのが一般的
セーブ
torch.save({
            'modelA_state_dict': modelA.state_dict(),
            'modelB_state_dict': modelB.state_dict(),
            'optimizerA_state_dict': optimizerA.state_dict(),
            'optimizerB_state_dict': optimizerB.state_dict(),
            ...
            }, PATH)
ロード
modelA = TheModelAClass(*args, **kwargs)
modelB = TheModelBClass(*args, **kwargs)
optimizerA = TheOptimizerAClass(*args, **kwargs)
optimizerB = TheOptimizerBClass(*args, **kwargs)

checkpoint = torch.load(PATH)
modelA.load_state_dict(checkpoint['modelA_state_dict'])
modelB.load_state_dict(checkpoint['modelB_state_dict'])
optimizerA.load_state_dict(checkpoint['optimizerA_state_dict'])
optimizerB.load_state_dict(checkpoint['optimizerB_state_dict'])

modelA.eval()
modelB.eval()
# - or -
modelA.train()
modelB.train()

他のモデルからパラメータをロードする場合

  • 転移学習をしたいときに便利
  • load_state_dict()strictFalseにすると,ロード側(modelB)のパラメータを設定する際,ロードされる側(modelA)にしか存在しないものを無視する
セーブ
torch.save(modelA.state_dict(), PATH)
ロード
modelB = TheModelBClass(*args, **kwargs)
modelB.load_state_dict(torch.load(PATH), strict=False)

デバイス間でモデルのセーブとロードを行う場合

GPUで訓練したモデルか,CPUで訓練したモデルかでロードのやり方が異なるので注意

GPUでセーブして,CPUでロードする場合
GPUでセーブ
torch.save(model.state_dict(), PATH)
CPUでロード
device = torch.device('cpu')
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))
GPUでセーブして,GPUでロードする場合
GPUでセーブ
torch.save(model.state_dict(), PATH)
GPUでロード
device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device)
CPUでセーブして,GPUでロードする場合
CPUでセーブ
torch.save(model.state_dict(), PATH)
GPUでロード
device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0"))  # Choose whatever GPU device number you want
model.to(device)
(おまけ)torch.nn.DataParallelモデルをセーブ/ロードする場合
セーブ
torch.save(model.module.state_dict(), PATH)
ロード

デバイスによって↑の通りに行う