ResNetで転移学習の方法を試してみる
この記事で学ぶこと
転移学習を簡単にまとめると、訓練済みのモデルがもつ特徴量を引き出す能力を利用する方法です。
例えば、ImageNetで訓練されたモデルを異なるデータセットへ適用して微調整をおこないます。
また、物体識別でバックボーン(特徴量を引きだす層)として利用したりする場合もあります。
今回はResNetを使って転移学習を行います。
ResNetをダウンロードする
まず、PipやCondaなどでPython環境を作ってください。バージョン３.８以降であれば問題ないです。PyTorchもインストールしておいてください。
MiDaSやYOLOv5のときと同様でTorch Hubを使いモデルをダウンロードします。
今回は、小さいモデルのResNet18を使います。
import torch
# pretrained=Trueで訓練済みのモデルがダウンロードされる。
model = torch.hub.load('pytorch/vision:v0.9.0', 'resnet18', pretrained=True)ダウンロードされたモデルは~/.cache/torch/hub/checkpoints/に格納されるので次回からは素早くロードできます。
ResNetをCIFAR10に適用する
ResNetはImageNetで訓練されたのですが、それを異なるデータセットであるCIFAR10で使えるようにします。
ImageNetは1000個のクラスがあるのに対し、CIFAR10は10個しかありません。
よってResNetの最後の層を1000クラスから10クラス対応に変更する必要があります。
まず、モデルをプリントして中身を見て見ましょう。
print(model)出力結果は長いので省略しますが、ResNetの特徴である残差ブロックが繰り返し出てきます。
今回のメインテーマは転移学習なので、モデル構造の詳細は省きます。
モデル自体に興味のある方はTorch Visionのソースコードあるいは論文を参照してください。
さて、モデルの最後に1000個の値を出すLinearレイヤーがあり、その名前がfcであるのがわかります。
(fc): Linear(in_features=512, out_features=1000, bias=True) 1000個の値がImageNetの1000個のクラスに対応しており、一番値の大きいクラスが画像識別の予測値となります。
CIFAR10の10個のクラスを予測するためにこのLinearレイヤーを変更します。
が、その前に、モデルの全てのパラメーターに対して勾配の計算をしないようにします。
# 勾配の計算をしないようにする
for param in model.parameters():
param.requires_grad = Falseこうすることで訓練済みのウェイトが変更されないようになりました。
次に、最後のLinearレイヤーを変更します。
model.fc = torch.nn.Linear(512, 10)再びモデルをプリントするとfcが変更されているのがわかります。
(fc): Linear(in_features=512, out_features=10, bias=True)この上書きされたLinearレイヤーのウェイトは勾配の計算がされるので、これから行う訓練によって調整されます。
これでモデルの変更が終わりました。
CIFAR10の訓練データをダウンロードする
以下は、PyTorchのこちらの記事を参考にしています。
CIFAR10のデータセットはTorch Visionから利用することができます。
詳しくはこちらのPyTorchのドキュメントを参照してください。
また、画像の前処理の仕方も上記の記事の解説に従ってTransformを作ります。
import torchvision
import torchvision.transforms as transforms
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
batch_size = 16
trainset = torchvision.datasets.CIFAR10(root='./data',
train=True,
download=True,
transform=transform)
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=batch_size,
shuffle=True,
num_workers=2)これでダウンロードできました。小さいデータセットなのですぐに終わります。
CIFAR10で訓練をする
訓練用に損失関数とオプティマイザーも設定します。
import torch.optim as optim
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.0005, momentum=0.9)これで訓練する準備はできました。
Linearレイヤーだけを訓練するので、下記のようにエポックの数を少なくしました。
# エポックの数
epochs = 2
# 訓練ループ
for epoch in range(epochs):
# ログ用にロスの値を格納する
running_loss = 0.0
# エポック
for i, (inputs, labels) in enumerate(trainloader, 0):
# 勾配(グラディエント)をゼロにする
optimizer.zero_grad()
# 予測値
outputs = model(inputs)
# 損失関数でロスを計算し、逆伝播する
loss = criterion(outputs, labels)
loss.backward()
# モデルの最適化調整
optimizer.step()
# print statistics
running_loss += loss.item()
# 200回に一回、平均のロスを表示
if (i+1) % 200 == 0:
print(f'[{epoch+1}, {i+1:5d}] loss: {running_loss/200:.3f}')
running_loss = 0.0
print('Finished Training')GPUがあるマシンを使っている方は、cudaデバイスを使うようにコードを修正すれば訓練が早く終わります。
最後にモデルのウェイトをセーブしておきます。
torch.save(model.state_dict(), 'resnet18.pt')CIFAR10でテストする
テスト用のデータは以下のようにロードします。
testset = torchvision.datasets.CIFAR10(root='./data',
train=False,
download=True,
transform=transform)
testloader = torch.utils.data.DataLoader(testset,
batch_size=batch_size,
shuffle=False,
num_workers=2)train=Falseとすることでテストデータを読み込みます。
また、shuffle=Falseに指定したのは、モデルの評価をする際に特にTrueにする必要がないからです。
モデルを評価モードにして、下記のように正解率を求めます。
# モデルを評価モードにする
model.eval()
# 正解と合計の数
correct = 0
total = 0
# テストループ
for i, (inputs, labels) in enumerate(testloader):
# 予測値
with torch.no_grad():
outputs = model(inputs)
# クラスの予測へと変換
prediction = outputs.argmax(axis=1)
# 正解の数を数える
correct += (labels==prediction).sum().item()
total += len(labels)
print(f'Finished Testing: accuracy={correct/total*100:.2f}')with torch.no_grad()を使っているのは、勾配の計算が必要ないからです。
余計な計算をしないことでプログラムが速く実行できます。
全部まとめたソース
import torch
from torch import nn, optim
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import CIFAR10
def train(model_name: str, epochs: int, batch_size: int, lr: float):
# pretrained=Trueで訓練済みのモデルがダウンロードされる。
model = torch.hub.load('pytorch/vision:v0.9.0', model_name, pretrained=True)
print(model)
# 勾配の計算をしないようにする
for param in model.parameters():
param.requires_grad = False
# モデルを10個のクラスの予測用に改良
model.fc = nn.Linear(512, 10)
print(model)
# データローダーを準備する
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
trainset = CIFAR10(root='./data',
train=True,
download=True,
transform=transform)
trainloader = DataLoader(trainset,
batch_size=batch_size,
shuffle=True,
num_workers=2)
# 損失関数とオプティマイザー
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)
# 訓練ループ
for epoch in range(epochs):
# ロス
running_loss = 0.0
# エポック
for i, (inputs, labels) in enumerate(trainloader):
# 勾配(グラディエント)をゼロにする
optimizer.zero_grad()
# 予測値
outputs = model(inputs)
# 損失関数でロスを計算し、逆伝播する
loss = criterion(outputs, labels)
loss.backward()
# モデルの最適化調整
optimizer.step()
# print statistics
running_loss += loss.item()
# 200回に一回、平均のロスを表示
if (i+1) % 200 == 0:
print(f'[{epoch+1}, {i+1:5d}] loss: {running_loss/200:.3f}')
running_loss = 0.0
torch.save(model.state_dict(), f'{model_name}.pt')
print('Finished Training')
def test(model_name: str, batch_size: int):
# pretrained=Trueで訓練済みのモデルがダウンロードされる。
model = torch.hub.load('pytorch/vision:v0.9.0', model_name, pretrained=True)
# モデルを10個のクラスの予測用に改良
model.fc = nn.Linear(512, 10)
# モデルのウェイトを読み込む
state_dict = torch.load(f'{model_name}.pt')
model.load_state_dict(state_dict)
# モデルを評価モードにする
model.eval()
# データローダーを準備する
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
testset = CIFAR10(root='./data',
train=False,
download=True,
transform=transform)
testloader = DataLoader(testset,
batch_size=batch_size,
shuffle=False,
num_workers=2)
# 正解と合計の数
correct = 0
total = 0
# テストループ
for i, (inputs, labels) in enumerate(testloader):
# 予測値
with torch.no_grad():
outputs = model(inputs)
# クラスの予測へと変換
prediction = outputs.argmax(axis=1)
# 正解の数を数える
correct += (labels==prediction).sum().item()
total += len(labels)
print(f'Finished Testing: accuracy={correct/total*100:.2f}')
if __name__=='__main__':
model_name = 'resnet18'
train(model_name, epochs=2, batch_size=16, lr=0.0005)
test(model_name, batch_size=32)
まとめ
今回の実験ではテストセットでの正解率は43.64％でした。
もっと高くできるはずですが、今回の本題は転移学習の方法の解説なので追求はしていません。
下記のように、いろいろとチューニングはできるとは思います。
エポックを増やす
バッチサイズを調節する
学習率を調整する（手動）
学習率を調整するスケジューラーを使ってみる
オプティマイザーを変えてみる（Adamとか）
訓練済みのレイヤーも微調整してみる（より小さい学習率を使う）
など、検討の余地はあります。やってみて下さい。
今日はこの辺で。
この記事が気に入ったらチップで応援してみませんか？