Posted at

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)


ロード

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