この記事では、YOLOv5のBackboneモジュールを実装する手順を説明します。
環境設定
- 言語: Python 3.8
- 開発環境: PyCharm
- データセット: 天気予測データセット（参照: 深度学習Day-03）
- ライブラリ: torch==1.12.1+cu113, torchvision==0.13.1+cu113
初期設定
1. GPUの設定
import torch
import torch.nn as nn
from torchvision import transforms, datasets
import warnings
warnings.filterwarnings("ignore") # 不要な警告を無視する
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)デバイスがGPU利用可能であればそれを使用し、それ以外はCPUを使用します。
2. データの読み込み
プロジェクトで使用するデータセットは公開されていませんので、ファイルディレクトリからデータを読み込みます。
import os, pathlib
data_dir = '../data'
data_dir = pathlib.Path(data_dir)
data_paths = list(data_dir.glob('*'))
class_names = [str(path).split(os.path.sep)[-1] for path in data_paths]
print(class_names)出力:
['cloudy', 'rain', 'shine', 'sunrise']次に、データセットに対する前処理を行います。
train_transforms = transforms.Compose([
transforms.Resize([224, 224]),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
test_transform = transforms.Compose([
transforms.Resize([224, 224]),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
total_data = datasets.ImageFolder("../data", transform=train_transforms)
print(total_data)出力:
Dataset ImageFolder
Number of datapoints: 1125
Root location: ../data
StandardTransform
Transform: Compose(
Resize(size=[224, 224], interpolation=bilinear, max_size=None, antialias=None)
RandomHorizontalFlip(p=0.5)
ToTensor()
Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
)クラスラベルをモデルが理解できる数値にマッピングします。
total_data.class_to_idx出力:
{'cloudy': 0, 'rain': 1, 'shine': 2, 'sunrise': 3}3. データセットの分割
データセットを訓練用とテスト用に分割します。
train_size = int(0.8 * len(total_data))
test_size = len(total_data) - train_size
train_dataset, test_dataset = torch.utils.data.random_split(total_data, [train_size, test_size])
batch_size = 4
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=0)
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=0)テストデータセットの形状を確認します。
for X, y in test_loader:
print("Shape of X [N, C, H, W]: ", X.shape)
print("Shape of y: ", y.shape, y.dtype)
break出力:
Shape of X [N, C, H, W]: torch.Size([4, 3, 224, 224])
Shape of y: torch.Size([4]) torch.int64Backboneモジュールを含むモデルの構築
1. モデルの定義
def pad_auto(k, p=None):
if p is None:
p = k // 2 if isinstance(k, int) else [x // 2 for x in k]
return p
class FeatureExtractor(nn.Module):
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):
super().__init__()
self.conv_layer = nn.Conv2d(c1, c2, k, s, pad_auto(k, p), groups=g, bias=False)
self.bn_layer = nn.BatchNorm2d(c2)
self.act_func = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
def forward(self, x):
return self.act_func(self.bn_layer(self.conv_layer(x)))
class ResidualBlock(nn.Module):
def __init__(self, c1, c2, shortcut=True, g=1, e=0.5):
super().__init__()
c_ = int(c2 * e)
self.conv1 = FeatureExtractor(c1, c_, 1, 1)
self.conv2 = FeatureExtractor(c_, c2, 3, 1, g=g)
self.shortcut = shortcut and c1 == c2
def forward(self, x):
return x + self.conv2(self.conv1(x)) if self.shortcut else self.conv2(self.conv1(x))
class CSPBlock(nn.Module):
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
super().__init__()
c_ = int(c2 * e)
self.conv1 = FeatureExtractor(c1, c_, 1, 1)
self.conv2 = FeatureExtractor(c1, c_, 1, 1)
self.conv3 = FeatureExtractor(2 * c_, c2, 1)
self.res_blocks = nn.Sequential(*(ResidualBlock(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
def forward(self, x):
return self.conv3(torch.cat((self.res_blocks(self.conv1(x)), self.conv2(x)), dim=1))
class SpatialPyramidPooling(nn.Module):
def __init__(self, c1, c2, k=5):
super().__init__()
c_ = c1 // 2
self.conv1 = FeatureExtractor(c1, c_, 1, 1)
self.conv2 = FeatureExtractor(c_ * 4, c2, 1, 1)
self.pool = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
def forward(self, x):
x = self.conv1(x)
with warnings.catch_warnings():
warnings.simplefilter('ignore')
y1 = self.pool(x)
y2 = self.pool(y1)
return self.conv2(torch.cat([x, y1, y2, self.pool(y2)], 1))
class YOLOv5Backbone(nn.Module):
def __init__(self):
super(YOLOv5Backbone, self).__init__()
self.conv1 = FeatureExtractor(3, 64, 3, 2, 2)
self.conv2 = FeatureExtractor(64, 128, 3, 2)
self.csp3 = CSPBlock(128, 128)
self.conv4 = FeatureExtractor(128, 256, 3, 2)
self.csp5 = CSPBlock(256, 256)
self.conv6 = FeatureExtractor(256, 512, 3, 2)
self.csp7 = CSPBlock(512, 512)
self.conv8 = FeatureExtractor(512, 1024, 3, 2)
self.csp9 = CSPBlock(1024, 1024)
self.sppf = SpatialPyramidPooling(1024, 1024, 5)
self.classifier = nn.Sequential(
nn.Linear(in_features=65536, out_features=100),
nn.ReLU(),
nn.Linear(in_features=100, out_features=4)
)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = self.csp3(x)
x = self.conv4(x)
x = self.csp5(x)
x = self.conv6(x)
x = self.csp7(x)
x = self.conv8(x)
x = self.csp9(x)
x = self.sppf(x)
x = torch.flatten(x, start_dim=1)
x = self.classifier(x)
return x
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")
model = YOLOv5Backbone().to(device)
print(model)2. モデル情報の確認
import torchsummary as summary
summary.summary(model, (3, 224, 224))モデルの訓練
1. 訓練関数の作成
def train_model(loader, model, criterion, optimizer):
size = len(loader.dataset)
batches = len(loader)
running_loss, running_corrects = 0.0, 0
model.train()
for inputs, labels in loader:
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
_, preds = torch.max(outputs, 1)
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
epoch_loss = running_loss / size
epoch_acc = running_corrects.double() / size
return epoch_acc, epoch_loss2. テスト関数の作成
def test_model(loader, model, criterion):
size = len(loader.dataset)
batches = len(loader)
running_loss, running_corrects = 0.0, 0
model.eval()
with torch.no_grad():
for inputs, labels in loader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels)
_, preds = torch.max(outputs, 1)
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
epoch_loss = running_loss / size
epoch_acc = running_corrects.double() / size
return epoch_acc, epoch_loss3. モデルの訓練
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.CrossEntropyLoss()
num_epochs = 20
best_acc = 0.0
best_model_wts = copy.deepcopy(model.state_dict())
train_losses, train_accuracies, test_losses, test_accuracies = [], [], [], []
for epoch in range(num_epochs):
train_acc, train_loss = train_model(train_loader, model, criterion, optimizer)
test_acc, test_loss = test_model(test_loader, model, criterion)
if test_acc > best_acc:
best_acc = test_acc
best_model_wts = copy.deepcopy(model.state_dict())
train_losses.append(train_loss)
train_accuracies.append(train_acc)
test_losses.append(test_loss)
test_accuracies.append(test_acc)
print(f'Epoch {epoch + 1}/{num_epochs}, '
f'Train Acc: {train_acc:.4f}, Train Loss: {train_loss:.4f}, '
f'Test Acc: {test_acc:.4f}, Test Loss: {test_loss:.4f}')
model.load_state_dict(best_model_wts)
torch.save(model.state_dict(), 'best_model.pth')結果の可視化
1. Loss & Accuracy
import matplotlib.pyplot as plt
plt.figure(figsize=(12, 3))
plt.subplot(1, 2, 1)
plt.plot(range(num_epochs), train_accuracies, label='Train Accuracy')
plt.plot(range(num_epochs), test_accuracies, label='Test Accuracy')
plt.legend(loc='lower right')
plt.title('Accuracy')
plt.subplot(1, 2, 2)
plt.plot(range(num_epochs), train_losses, label='Train Loss')
plt.plot(range(num_epochs), test_losses, label='Test Loss')
plt.legend(loc='upper right')
plt.title('Loss')
plt.show()2. モデル評価
model.eval()
test_acc, test_loss = test_model(test_loader, model, criterion)
print(f'Test Accuracy: {test_acc:.4f}, Test Loss: {test_loss:.4f}')