今回は自作のデータセットを読み込み、PyTorchの物体検出モデルで学習(ファインチューニング)させる手順についてまとめました。
自作のデータセットにはLabelmeで作成したアノテーションデータを使用し、物体検出モデルにはTorchVisionに実装されているFaster R-CNNを使用します。
※手順をまとめることが目的なので実用的なデータは使用していません。
1. 前処理の追加
まず初めに、前処理として「リサイズ」と「データ水増し」と呼ばれる処理を追加していきます。
水増しとは、画像を回転させたりして1枚の画像から複数の異なるデータを生成し、疑似的にデータを増やすことにより精度向上が見込めます。
今回は例としてリサイズ、左右反転、上下反転、明るさ変更、ランダムノイズ追加処理を実装します。
1.1. リサイズ
リサイズ処理を行うクラスは以下の通りです。
probが1.0以外の時に画像とバウンティングボックスそれぞれに対してリサイズ処理を実行しています。
# リサイズ
class Trans_Resize(object):
def __init__(self, prob):
self.prob = prob
def __call__(self, image, target):
c, w, h = image.shape
if self.prob != 1.0:
image = F.resize(img=image, size=(int(w*self.prob), int(h*self.prob)))
target["boxes"][:,:] = target["boxes"][:,:] * self.prob
return image, target1.2. 左右反転
左右反転処理を行うクラスは以下の通りです。
probが乱数より大きい場合に画像とバウンティングボックスそれぞれに対して左右反転処理を実行します。
#左右反転
class RandomHorizontalFlip(object):
def __init__(self, prob):
self.prob = prob
def __call__(self, image, target):
if random.random() < self.prob:
w = image.shape[2]
image = image.flip(2)
bbox = target["boxes"]
bbox[:, [0, 2]] = w - bbox[:, [2, 0]]
target["boxes"] = bbox
return image, target1.3. 上下反転
上下反転処理を行うクラスは以下の通りです。
左右反転処理と同様、probが乱数より大きい場合に画像とバウンティングボックスそれぞれに対して上下反転処理を実行します。
#上下反転
class RandomVarticalFlip(object):
def __init__(self, prob):
self.prob = prob
def __call__(self, image, target):
if random.random() < self.prob:
h = image.shape[1]
image = image.flip(1)
bbox = target["boxes"]
bbox[:, [1, 3]] = h - bbox[:, [3, 1]]
target["boxes"] = bbox
return image, target1.4. 明るさ変更(ガンマ補正)
明るさの変更を行うクラスは以下の通りです。
probが乱数より大きい場合に、0.5～1.5の範囲でGamma値を生成して画像に適用します。
なおバウンティングボックスは変更がないのでtargetはそのまま返します。
#明るさ変更(ガンマ補正)
class RandomBrightness(object):
def __init__(self, prob):
self.prob = prob
def __call__(self, image, target):
if random.random() < self.prob:
gamma = random.randint(5, 15) * 0.1
image[:,:,:] = pow(image[:,:,:], 1.0/gamma)
print(gamma)
return image, target1.5. ランダムノイズ追加
ランダムノイズを行うクラスは以下の通りです。
probがが乱数より大きい場合に、画像サイズの1/8～1/4の数だけランダムな位置に点を打ちます。
なお明るさ同様バウンティングボックスは変更がないのでtargetはそのまま返します。
#ランダムノイズ
class RandomNoise(object):
def __init__(self, prob):
self.prob = prob
def __call__(self, image, target):
if random.random() < self.prob:
height, width = image.shape[-2:]
pts_count = np.random.randint((height*width)/8, (height*width)/4, 1)
r = torch.as_tensor(np.random.rand(pts_count.item()), dtype=torch.float32)
g = torch.as_tensor(np.random.rand(pts_count.item()), dtype=torch.float32)
b = torch.as_tensor(np.random.rand(pts_count.item()), dtype=torch.float32)
pts_rgb = torch.reshape(torch.cat([r,g,b],dim=0), (3,pts_count.item()))
pts_x = np.random.randint(0, width-1 , pts_count.item())
pts_y = np.random.randint(0, height-1, pts_count.item())
image[:,pts_y,pts_x] = pts_rgb
return image, target1.6. 前処理用関数の実装
上記で実装したクラスを加えた前処理用関数を「get_transform」として実装します。
import torch
import random
import numpy as np
from torchvision.transforms import functional as F
class ToTensor(object):
def __call__(self, image, target):
image = F.to_tensor(image)
return image, target
# リサイズ
class Trans_Resize(object):
def __init__(self, prob):
self.prob = prob
def __call__(self, image, target):
c, w, h = image.shape
if self.prob != 1.0:
image = F.resize(img=image, size=(int(w*self.prob), int(h*self.prob)))
target["boxes"][:,:] = target["boxes"][:,:] * self.prob
return image, target
#左右反転
class RandomHorizontalFlip(object):
def __init__(self, prob):
self.prob = prob
def __call__(self, image, target):
if random.random() < self.prob:
w = image.shape[2]
image = image.flip(2)
bbox = target["boxes"]
bbox[:, [0, 2]] = w - bbox[:, [2, 0]]
target["boxes"] = bbox
return image, target
#上下反転
class RandomVarticalFlip(object):
def __init__(self, prob):
self.prob = prob
def __call__(self, image, target):
if random.random() < self.prob:
h = image.shape[1]
image = image.flip(1)
bbox = target["boxes"]
bbox[:, [1, 3]] = h - bbox[:, [3, 1]]
target["boxes"] = bbox
return image, target
#明るさ変更(ガンマ補正)
class RandomBrightness(object):
def __init__(self, prob):
self.prob = prob
def __call__(self, image, target):
if random.random() < self.prob:
gamma = random.randint(5, 15) * 0.1
image[:,:,:] = pow(image[:,:,:], 1.0/gamma)
return image, target
#ランダムノイズ
class RandomNoise(object):
def __init__(self, prob):
self.prob = prob
def __call__(self, image, target):
if random.random() < self.prob:
height, width = image.shape[-2:]
pts_count = np.random.randint((height*width)/8, (height*width)/4, 1)
r = torch.as_tensor(np.random.rand(pts_count.item()), dtype=torch.float32)
g = torch.as_tensor(np.random.rand(pts_count.item()), dtype=torch.float32)
b = torch.as_tensor(np.random.rand(pts_count.item()), dtype=torch.float32)
pts_rgb = torch.reshape(torch.cat([r,g,b],dim=0), (3,pts_count.item()))
pts_x = np.random.randint(0, width-1 , pts_count.item())
pts_y = np.random.randint(0, height-1, pts_count.item())
image[:,pts_y,pts_x] = pts_rgb
return image, target
class Compose(object):
def __init__(self, transforms):
self.transforms = transforms
def __call__(self, image, target):
for t in self.transforms:
image, target = t(image, target)
return image, target
#前処理関数
def get_transform(train=False,
resize = 1.0,
hflip = 0.0,
vflip = 0.0,
brightness = 0.0,
noise = 0.0):
transforms = []
# PIL imageをPyTorch Tensorに変換
transforms.append(ToTensor())
transforms.append(Trans_Resize(resize))
if train:
#学習時のみの処理を追加
transforms.append(RandomHorizontalFlip(hflip))
transforms.append(RandomVarticalFlip(vflip))
transforms.append(RandomBrightness(brightness))
transforms.append(RandomNoise(noise))
return Compose(transforms)2. データセット
2.1. データセットクラス
次にデータセットの読み込み部分を実装します。
今回使用するデータセット読み込み用クラスは以下の通りです（以前の記事で作成したカスタムデータセットと同じ）。
import torch
import json
from PIL import Image
import base64
import glob
import io
import copy
class Custom_Dataset(torch.utils.data.Dataset):
def __init__(self, root, transforms):
self.imgs = []
self.targets = []
self.transforms = transforms
self.CreateDataset(root)
def __len__(self):
return len(self.imgs)
def __getitem__(self, idx):
img = self.imgs[idx]
target = copy.deepcopy(self.targets[idx])
if self.transforms is not None:
img, target = self.transforms(img, target)
return img, target
# base64形式をPIL型に変換
def img_data_to_pil(self,img_data):
f = io.BytesIO()
f.write(img_data)
img_pil = Image.open(f)
return img_pil
def CreateDataset(self, json_dir):
# jsonデータの読み込み
json_paths = glob.glob(json_dir + '/*.json')
for json_path in json_paths:
json_file = open(json_path)
json_data = json.load(json_file)
# imageDataをkeyにしてデータを取り出す
img_b64 = json_data['imageData']
#base64形式をPIL型に変換する
img_data = base64.b64decode(img_b64)
img_pil = self.img_data_to_pil(img_data)
# imgsに画像を追加
self.imgs.append(img_pil)
# targets作成
num_objs = len(json_data['shapes']) # 物体の個数
boxes = []
# バウンティングボックスから左上、左下、右上、右下の座標を取得し、
# boxesリストへ追加
for i in range(num_objs):
box = json_data['shapes'][i]['points']
x_list = []
y_list = []
for i in range(len(box)):
x,y=box[i]
x_list.append(x)
y_list.append(y)
xmin = min(x_list)
xmax = max(x_list)
ymin = min(y_list)
ymax = max(y_list)
boxes.append([xmin, ymin, xmax, ymax])
# boxesリストをtorch.Tensorに変換
boxes = torch.as_tensor(boxes, dtype=torch.float32)
# torch.Tensorにラベルを設定
# 物体の個数(背景を0とする)
labels = torch.ones((num_objs,), dtype=torch.int64)
#targetsへboxesリストとlabelsリストを設定
target = {}
target["boxes"] = boxes
target["labels"] = labels
self.targets.append(target)画像を1枚表示してみます。
import torchvision
from PIL import ImageDraw
#前処理
from transforms import get_transform
#カスタムデータセット
from CustomDataset import Custom_Dataset
dataset = Custom_Dataset('dataset',get_transform())
img, target = dataset[0]
# Tensor -> PILへ変換
img = torchvision.transforms.functional.to_pil_image(img)
# バウンティングボックス描画
for i in range(len(target['boxes'])):
draw = ImageDraw.Draw(img)
x = target['boxes'][i][0]
y = target['boxes'][i][1]
w = target['boxes'][i][2]
h = target['boxes'][i][3]
draw.rectangle((x,y,w,h), outline=(255, 0, 0), width=3)
img.show()実行すると以下のように画像が表示されバウンティングボックスが描画されます。
2.2. 前処理の確認
ここで一旦前処理の確認をしてみます。
import torchvision
from PIL import ImageDraw
from transforms import get_transform
from CustomDataset import Custom_Dataset
from transforms import Trans_Resize, RandomHorizontalFlip, RandomVarticalFlip, RandomBrightness, RandomNoise
dataset = Custom_Dataset('dataset',get_transform())
trans_proc = [Trans_Resize(0.5), RandomHorizontalFlip(1.0), RandomVarticalFlip(1.0), RandomBrightness(1.0), RandomNoise(1.0)]
for j in range(len(trans_proc)):
img, target = dataset[0]
img, target = trans_proc[j](img, target)
img = torchvision.transforms.functional.to_pil_image(img)
# バウンティングボックス描画
for i in range(len(target['boxes'])):
draw = ImageDraw.Draw(img)
x = target['boxes'][i][0]
y = target['boxes'][i][1]
w = target['boxes'][i][2]
h = target['boxes'][i][3]
draw.rectangle((x,y,w,h), outline=(255, 0, 0), width=3)
img.show()結果はそれぞれ以下の通りです。
・リサイズ
・左右反転
・上下反転
・明るさ変更(Gamma=0.5)
・ランダムノイズ追加
正常に処理されていることが確認できました。
3. データローダーの定義
次にデータローダーを定義します。データローダーはデータセットを入力することでバッチサイズごとの取り出し、 データのシャッフル、並列処理などを自動で行ってくれます。
import torch
import torchvision
from PIL import ImageDraw
from transforms import get_transform
from CustomDataset import Custom_Dataset
def collate_fn(batch):
return tuple(zip(*batch))
dataset = Custom_Dataset('dataset',get_transform())
#バッチサイズ
batch_size = 1
#データローダー定義
data_loader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size, shuffle=True, num_workers=0,
collate_fn=collate_fn,pin_memory=True)collate_fnはデータの出力形式を指定する関数です。imgとtargetはリスト形式になっているので、それをバッチサイズごとにtuple形式に変換しています。
また、データローダーから取り出されるデータはミニバッチごとにまとまっているので1次元追加されます。
dataset : img=[3, 512, 512] → data_loader : imgs=[1, 3, 512, 512]
4. モデルの読み込み
次に物体検出モデルの読み込みを行います。今回はTorchVisionに実装されている「Faster R-CNN」を使用します。
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
def get_instance_model(num_classes):
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes=num_classes)
return modelそして以下のようにモデルを読み込みを行います。検出したい物体が1つの場合は、背景と合わせてnum_classesは2に設定します。
num_classes = 2
net = model.get_instance_model(num_classes)5. デバイス設定
次にデバイスの設定を行います。GPUが使用可能な場合は「cuda」、使用できない場合は「cpu」が設定されます。
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
net.to(device)6. オプティマイザとスケジューラの定義
次にオプティマイザとスケジューラを定義します。
「オプティマイザ」ではlossを収束させるための最適化関数を選択します。ここではSGDを選択します。
「スケジューラ」を使用することで、学習途中で学習率(lr)を変更することができます。ここではStepLRを選択します。
params = [p for p in net.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005,momentum=0.9, weight_decay=0.0005)
#3エポックごとにlrを0.5倍
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,step_size=3,gamma=0.5)7. 学習用関数定義
次に学習用の関数を定義します。
import torch
def train(model, optimizer, data_loader, device):
loss_ave = 0
model.train()
for imgs, targets in data_loader:
# Faster R-CNNは入力がList形式なので
# img, targetそれぞれ変換
img = list(image.to(device, dtype=torch.float) for image in imgs)
target = [{k: v.to(device) for k, v in t.items()} for t in targets]
# imgとtargetをモデルに入力し、lossを計算
loss_dict = model(img, target)
# 4種類のlossが出力されるので合計を算出
losses = sum(loss for loss in loss_dict.values())
print(loss_dict)
print(losses)
loss_value = losses.item()
loss_ave += loss_value
#勾配リセット
optimizer.zero_grad()
#バックプロパゲーション算出
losses.backward()
#重み更新
optimizer.step()
# Lossの平均を返す
return loss_ave / len(data_loader)そして以下のように学習処理を実装します。
#エポック数
num_epochs = 10
loss = 0
loss_max = 0
loss_min = 0
epoch_min = 0
start_time = time.time()
for epoch in range(num_epochs):
loss = train(net, optimizer, data_loader, device, epoch)
lr=optimizer.param_groups[0]["lr"]
if epoch == 0:
loss_max = loss
loss_min = loss
else :
if loss_max < loss:
loss_max = loss
elif loss_min > loss:
loss_min = loss
epoch_min = epoch
# モデルデータ保存
torch.save(model, 'test.pth')
end_time = time.time()
print("Epoch[{}] 平均Loss [{}], 最小Loss[{}], 最小Loss_epoch[{}], 学習率 [{}, Total_Time[{}]]".format(epoch, loss, loss_min, epoch_min, lr, end_time - start_time))
# 学習率の更新
lr_scheduler.step()
print("処理時間:{0}".format(end_time - start_time))8. プログラム全体
プログラム全体は以下の通りです。
import torch
import json
from PIL import Image
import base64
import glob
import io
import copy
class Custom_Dataset(torch.utils.data.Dataset):
def __init__(self, root, transforms):
self.imgs = []
self.targets = []
self.transforms = transforms
self.CreateDataset(root)
def __len__(self):
return len(self.imgs)
def __getitem__(self, idx):
img = self.imgs[idx]
target = copy.deepcopy(self.targets[idx])
if self.transforms is not None:
img, target = self.transforms(img, target)
return img, target
# base64形式をPIL型に変換
def img_data_to_pil(self,img_data):
f = io.BytesIO()
f.write(img_data)
img_pil = Image.open(f)
return img_pil
def CreateDataset(self, json_dir):
# jsonデータの読み込み
json_paths = glob.glob(json_dir + '/*.json')
for json_path in json_paths:
json_file = open(json_path)
json_data = json.load(json_file)
# imageDataをkeyにしてデータを取り出す
img_b64 = json_data['imageData']
#base64形式をPIL型に変換する
img_data = base64.b64decode(img_b64)
img_pil = self.img_data_to_pil(img_data)
# imgsに画像を追加
self.imgs.append(img_pil)
# targets作成
num_objs = len(json_data['shapes']) # 物体の個数
boxes = []
# バウンティングボックスから左上、左下、右上、右下の座標を取得し、
# boxesリストへ追加
for i in range(num_objs):
box = json_data['shapes'][i]['points']
x_list = []
y_list = []
for i in range(len(box)):
x,y=box[i]
x_list.append(x)
y_list.append(y)
xmin = min(x_list)
xmax = max(x_list)
ymin = min(y_list)
ymax = max(y_list)
boxes.append([xmin, ymin, xmax, ymax])
# boxesリストをtorch.Tensorに変換
boxes = torch.as_tensor(boxes, dtype=torch.float32)
# torch.Tensorにラベルを設定
# 物体の個数(背景を0とする)
labels = torch.ones((num_objs,), dtype=torch.int64)
#targetsへboxesリストとlabelsリストを設定
target = {}
target["boxes"] = boxes
target["labels"] = labels
self.targets.append(target)import torch
import random
import numpy as np
from torchvision.transforms import functional as F
class ToTensor(object):
def __call__(self, image, target):
image = F.to_tensor(image)
return image, target
# リサイズ
class Trans_Resize(object):
def __init__(self, prob):
self.prob = prob
def __call__(self, image, target):
c, w, h = image.shape
if self.prob != 1.0:
image = F.resize(img=image, size=(int(w*self.prob), int(h*self.prob)))
target["boxes"][:,:] = target["boxes"][:,:] * self.prob
return image, target
#左右反転
class RandomHorizontalFlip(object):
def __init__(self, prob):
self.prob = prob
def __call__(self, image, target):
if random.random() < self.prob:
w = image.shape[2]
image = image.flip(2)
bbox = target["boxes"]
bbox[:, [0, 2]] = w - bbox[:, [2, 0]]
target["boxes"] = bbox
return image, target
#上下反転
class RandomVarticalFlip(object):
def __init__(self, prob):
self.prob = prob
def __call__(self, image, target):
if random.random() < self.prob:
h = image.shape[1]
image = image.flip(1)
bbox = target["boxes"]
bbox[:, [1, 3]] = h - bbox[:, [3, 1]]
target["boxes"] = bbox
return image, target
#明るさ変更(ガンマ補正)
class RandomBrightness(object):
def __init__(self, prob):
self.prob = prob
def __call__(self, image, target):
if random.random() < self.prob:
gamma = random.randint(5, 15) * 0.1
image[:,:,:] = pow(image[:,:,:], 1.0/gamma)
return image, target
#ランダムノイズ
class RandomNoise(object):
def __init__(self, prob):
self.prob = prob
def __call__(self, image, target):
if random.random() < self.prob:
height, width = image.shape[-2:]
pts_count = np.random.randint((height*width)/8, (height*width)/4, 1)
r = torch.as_tensor(np.random.rand(pts_count.item()), dtype=torch.float32)
g = torch.as_tensor(np.random.rand(pts_count.item()), dtype=torch.float32)
b = torch.as_tensor(np.random.rand(pts_count.item()), dtype=torch.float32)
pts_rgb = torch.reshape(torch.cat([r,g,b],dim=0), (3,pts_count.item()))
pts_x = np.random.randint(0, width-1 , pts_count.item())
pts_y = np.random.randint(0, height-1, pts_count.item())
image[:,pts_y,pts_x] = pts_rgb
return image, target
class Compose(object):
def __init__(self, transforms):
self.transforms = transforms
def __call__(self, image, target):
for t in self.transforms:
image, target = t(image, target)
return image, target
def get_transform(train=False,
resize = 1.0,
hflip = 0.0,
vflip = 0.0,
brightness = 0.0,
noise = 0.0):
transforms = []
# PIL imageをPyTorch Tensorに変換
transforms.append(ToTensor())
transforms.append(Trans_Resize(resize))
if train:
#学習時のみの処理を追加
transforms.append(RandomHorizontalFlip(hflip))
transforms.append(RandomVarticalFlip(vflip))
transforms.append(RandomBrightness(brightness))
transforms.append(RandomNoise(noise))
return Compose(transforms)import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
def get_instance_model(num_classes):
# COCOデータセットで訓練した、訓練済みモデルをロード
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
# 既存の分類器を、ユーザーが定義したnum_classesを持つ新しい分類器に置き換えます
# 分類器にインプットする特徴量の数を取得
in_features = model.roi_heads.box_predictor.cls_score.in_features
# 事前訓練済みのヘッドを新しいものと置き換える
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes=num_classes)
return modelimport torch
def train(model, optimizer, data_loader, device):
loss_ave = 0
model.train()
for imgs, targets in data_loader:
# Faster R-CNNは入力がList形式なので
# img, targetそれぞれ変換
img = list(image.to(device, dtype=torch.float) for image in imgs)
target = [{k: v.to(device) for k, v in t.items()} for t in targets]
# imgとtargetをモデルに入力し、lossを計算
loss_dict = model(img, target)
# 4種類のlossが出力されるので合計を算出
losses = sum(loss for loss in loss_dict.values())
loss_value = losses.item()
loss_ave += loss_value
#勾配リセット
optimizer.zero_grad()
#バックプロパゲーション算出
losses.backward()
#重み更新
optimizer.step()
# Lossの平均を返す
return loss_ave / len(data_loader)import torch
import time
import sys
from CustomDataset import Custom_Dataset
import model
from train import train
from transforms import get_transform
def collate_fn(batch):
return tuple(zip(*batch))
if __name__ == '__main__':
dataset = Custom_Dataset('dataset',get_transform(train=True, resize=1.0, hflip=0.5, vflip=0.5,brightness=0.5,noise=0.5))
#バッチサイズ
batch_size = 1
#データローダー定義
data_loader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size, shuffle=True, num_workers=0,
collate_fn=collate_fn,pin_memory=True)
num_classes = 2
net = model.get_instance_model(num_classes)
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
print(device)
net.to(device)
params = [p for p in net.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005,momentum=0.9, weight_decay=0.0005)
#3エポックごとにlrを0.5倍
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,step_size=3,gamma=0.5)
#エポック数
num_epochs = 10
loss = 0
loss_max = 0
loss_min = 0
epoch_min = 0
start_time = time.time()
for epoch in range(num_epochs):
loss = train(net, optimizer, data_loader, device)
lr=optimizer.param_groups[0]["lr"]
if epoch == 0:
loss_max = loss
loss_min = loss
else :
if loss_max < loss:
loss_max = loss
elif loss_min > loss:
loss_min = loss
epoch_min = epoch
# モデルデータ保存
torch.save(net, 'test.pth')
end_time = time.time()
print("Epoch[{}] 平均Loss [{}], 最小Loss[{}], 最小Loss_epoch[{}], 学習率 [{}, Total_Time[{}]]".format(epoch, loss, loss_min, epoch_min, lr, end_time - start_time))
# 学習率の更新
lr_scheduler.step()
print("処理時間:{0}".format(end_time - start_time))9. 学習処理実行
実際に学習処理を動かしてみました。動かす際はGoogle Colaboratoryを使用したのでmain.pyを以下のように修正しました。
とりあえず20エポックで。
import torch
import time
import sys
sys.path.append('/content/drive/My Drive/Colab Notebooks/test')
from CustomDataset import Custom_Dataset
import model
from train import train
from transforms import get_transform
def collate_fn(batch):
return tuple(zip(*batch))
if __name__ == '__main__':
dataset = Custom_Dataset('/content/drive/My Drive/Colab Notebooks/test/dataset',get_transform(train=True, resize=1.0, hflip=0.5, vflip=0.5,brightness=0.5,noise=0.5))
#バッチサイズ
batch_size = 1
#データローダー定義
data_loader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size, shuffle=True, num_workers=0,
collate_fn=collate_fn,pin_memory=True)
num_classes = 2
net = model.get_instance_model(num_classes)
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
print(device)
net.to(device)
params = [p for p in net.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005,momentum=0.9, weight_decay=0.0005)
#3エポックごとにlrを0.5倍
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,step_size=3,gamma=0.5)
#エポック数
num_epochs = 10
loss = 0
loss_max = 0
loss_min = 0
epoch_min = 0
start_time = time.time()
for epoch in range(num_epochs):
loss = train(net, optimizer, data_loader, device)
lr=optimizer.param_groups[0]["lr"]
if epoch == 0:
loss_max = loss
loss_min = loss
else :
if loss_max < loss:
loss_max = loss
elif loss_min > loss:
loss_min = loss
epoch_min = epoch
# モデルデータ保存
torch.save(net, '/content/drive/My Drive/Colab Notebooks/test/test.pth')
end_time = time.time()
print("Epoch[{}] 平均Loss [{}], 最小Loss[{}], 最小Loss_epoch[{}], 学習率 [{}, Total_Time[{}]]".format(epoch, loss, loss_min, epoch_min, lr, end_time - start_time))
# 学習率の更新
lr_scheduler.step()
print("処理時間:{0}".format(end_time - start_time))
cuda
Epoch[0] 平均Loss [0.17168771581990377], 最小Loss[0.17168771581990377], 最小Loss_epoch[0], 学習率 [0.005, Total_Time[5.391830205917358]]
Epoch[1] 平均Loss [0.09176834566252572], 最小Loss[0.09176834566252572], 最小Loss_epoch[1], 学習率 [0.005, Total_Time[10.664321660995483]]
Epoch[2] 平均Loss [0.06550697166295279], 最小Loss[0.06550697166295279], 最小Loss_epoch[2], 学習率 [0.005, Total_Time[16.17812705039978]]
Epoch[3] 平均Loss [0.040696475033958755], 最小Loss[0.040696475033958755], 最小Loss_epoch[3], 学習率 [0.0025, Total_Time[21.509663820266724]]
Epoch[4] 平均Loss [0.035694786657889686], 最小Loss[0.035694786657889686], 最小Loss_epoch[4], 学習率 [0.0025, Total_Time[27.415638208389282]]
Epoch[5] 平均Loss [0.0321974911327873], 最小Loss[0.0321974911327873], 最小Loss_epoch[5], 学習率 [0.0025, Total_Time[32.82360053062439]]
Epoch[6] 平均Loss [0.025488603372304214], 最小Loss[0.025488603372304214], 最小Loss_epoch[6], 学習率 [0.00125, Total_Time[38.24358606338501]]
Epoch[7] 平均Loss [0.02130358927838859], 最小Loss[0.02130358927838859], 最小Loss_epoch[7], 学習率 [0.00125, Total_Time[43.83955907821655]]
Epoch[8] 平均Loss [0.019849855585822036], 最小Loss[0.019849855585822036], 最小Loss_epoch[8], 学習率 [0.00125, Total_Time[49.276163816452026]]
Epoch[9] 平均Loss [0.017080800474754403], 最小Loss[0.017080800474754403], 最小Loss_epoch[9], 学習率 [0.000625, Total_Time[54.66871905326843]]
Epoch[10] 平均Loss [0.015477336987498262], 最小Loss[0.015477336987498262], 最小Loss_epoch[10], 学習率 [0.000625, Total_Time[60.04023480415344]]
Epoch[11] 平均Loss [0.01650084253577959], 最小Loss[0.015477336987498262], 最小Loss_epoch[10], 学習率 [0.000625, Total_Time[65.34068608283997]]
Epoch[12] 平均Loss [0.012868352788722231], 最小Loss[0.012868352788722231], 最小Loss_epoch[12], 学習率 [0.0003125, Total_Time[70.76740264892578]]
Epoch[13] 平均Loss [0.013304523746704771], 最小Loss[0.012868352788722231], 最小Loss_epoch[12], 学習率 [0.0003125, Total_Time[76.08505606651306]]
Epoch[14] 平均Loss [0.012885955228869404], 最小Loss[0.012868352788722231], 最小Loss_epoch[12], 学習率 [0.0003125, Total_Time[81.60282278060913]]
Epoch[15] 平均Loss [0.013484267384878226], 最小Loss[0.012868352788722231], 最小Loss_epoch[12], 学習率 [0.00015625, Total_Time[87.054283618927]]
Epoch[16] 平均Loss [0.01088811204369579], 最小Loss[0.01088811204369579], 最小Loss_epoch[16], 学習率 [0.00015625, Total_Time[92.35550427436829]]
Epoch[17] 平均Loss [0.010897022494602771], 最小Loss[0.01088811204369579], 最小Loss_epoch[16], 学習率 [0.00015625, Total_Time[97.83080339431763]]
Epoch[18] 平均Loss [0.013651401735842228], 最小Loss[0.01088811204369579], 最小Loss_epoch[16], 学習率 [7.8125e-05, Total_Time[103.17689347267151]]
Epoch[19] 平均Loss [0.010561094042800721], 最小Loss[0.010561094042800721], 最小Loss_epoch[19], 学習率 [7.8125e-05, Total_Time[108.51215744018555]]
処理時間:108.51215744018555最後に、学習したモデルをCPUのPCで読み込み、検出結果を表示してみる。
import torchvision
import torch
from PIL import ImageDraw
from CustomDataset import Custom_Dataset
from transforms import get_transform
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
dataset = Custom_Dataset('dataset',get_transform())
img, _ = dataset[0]
#GPUで学習したモデルをCPUで読み込み場合は「map_loacation」をcpuに変更
net = torch.load('test.pth', map_location=torch.device('cpu'))
net.to(device)
# モデルを評価モードに変更します
net.eval()
with torch.no_grad():
prediction = net([img.to(device)])
print(len(prediction))
print(prediction[0]["boxes"][0])
img = torchvision.transforms.functional.to_pil_image(img)
# バウンティングボックス描画
for i in range(len(prediction[0]["boxes"])):
draw = ImageDraw.Draw(img)
x = prediction[0]['boxes'][i][0]
y = prediction[0]['boxes'][i][1]
w = prediction[0]['boxes'][i][2]
h = prediction[0]['boxes'][i][3]
draw.rectangle((x,y,w,h), outline=(255, 0, 0), width=3)
img.show() 結果はこちら。
正常に動作することが確認できました。
今回は以上です。
10. 参考文献・参考サイト
書籍
・赤石 雅典 ,日経BP,「最短コースでわかる PyTorch &深層学習プログラミング」,2021.09.21
https://bookplus.nikkei.com/atcl/catalog/21/283980/
Webサイト
・「Torchvisionを利用した物体検出のファインチューニング手法」
https://colab.research.google.com/github/YutaroOgawa/pytorch_tutorials_jp/blob/main/notebook/2_Image_Video/2_2_torchvision_finetuning_instance_segmentation_jp.ipynb
・PyTorchの高速化テク！pin_memoryとnum_workersでデータローディングを最適化しよう！
https://qiita.com/mizomizo1/items/d89d099823af690fbbeb
・たった1行でデータ拡張！torchvision のコードまとめ！
https://www.kikagaku.co.jp/kikagaku-blog/pytorch-torchvision/#i-26
・Pytorchのcollate_fnを使ってみる
https://qiita.com/tomp/items/f220bd6ffec006dabaa5
・DataLoaderの中身の確認: for文を使う方法とiterとnextを 使う方法
https://output-zakki.com/dataloader_iter_and_next/
・【前編】Pytorchの様々な最適化手法(torch.optim.Optimizer)の更新過程や性能を比較検証してみた！
https://rightcode.co.jp/blogs/17100
・【PyTorch】学習率スケジューラー解説
https://zenn.dev/yuto_mo/articles/6e2803495029d4
コメント