過去にKerasを利用した画像2クラス分類の転移学習コードを公開しましたが、開発をPytorchに移行したため、いつ・どこからでも参照できるようにブログ上にPytorch版を公開しておきます。
本ページではVGG16を利用した転移学習による2クラス分類モデルの実装を扱います。
開発環境
python: 3.12.11
matplotlib: 3.10.6
tqdm: 4.67.1
pytorch: 2.7.1+cu118
ソースコード
from typing import List, Tuple
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, models, transforms
from tqdm.auto import tqdm, trange # 追加
# ====== 1) データ読み込み：CIFAR-10のクラス0/1のみ抽出 ======
def get_cifar10_two_classes(
root: str = "./data",
classes_keep: List[int] = [0, 1],
) -> Tuple[Subset, Subset]:
"""
CIFAR-10 からクラス0/1のみを抽出した train/test の Subset を返します。
"""
train_transform = transforms.Compose(
[
transforms.Resize((32, 32)), # 念のため固定
transforms.RandomRotation(degrees=90),
transforms.RandomAffine(degrees=0, translate=(0.2, 0.2)),
transforms.RandomHorizontalFlip(p=0.5),
transforms.ToTensor(), # [0,1]スケーリング
# NOTE: ここでNormalizeを入れるとImageNet学習済みと相性◎（後述）
# transforms.Normalize(mean=[0.485,0.456,0.406],
# std=[0.229,0.224,0.225]),
]
)
test_transform = transforms.Compose(
[
transforms.Resize((32, 32)),
transforms.ToTensor(),
]
)
trainset = datasets.CIFAR10(
root=root, train=True, download=True, transform=train_transform
)
testset = datasets.CIFAR10(
root=root, train=False, download=True, transform=test_transform
)
# クラス0/1のみのインデックスを抽出
def filter_indices(ds: datasets.CIFAR10, keep: List[int]) -> List[int]:
# datasets.CIFAR10.targets は list[int]
return [i for i, t in enumerate(ds.targets) if t in keep]
train_idx = filter_indices(trainset, classes_keep)
test_idx = filter_indices(testset, classes_keep)
return Subset(trainset, train_idx), Subset(testset, test_idx)
# ====== 2) モデル定義：VGG16の畳み込み部を凍結し、ヘッドだけ学習 ======
class VGG16BinaryClassifier(nn.Module):
"""
- features（畳み込み部）: ImageNet学習済みを利用、学習は凍結
- head: AdaptiveAvgPool2d(1) -> Flatten -> Linear(512->1)
出力はロジット（BCEWithLogitsLossを使用）
"""
def __init__(self, pretrained: bool = True, freeze_features: bool = True):
super().__init__()
base = models.vgg16(
weights=models.VGG16_Weights.IMAGENET1K_V1 if pretrained else None
)
self.features = base.features # 畳み込み部
if freeze_features:
for p in self.features.parameters():
p.requires_grad = False
self.pool = nn.AdaptiveAvgPool2d((1, 1))
self.head = nn.Linear(512, 1) # 2値分類: ロジット1本（sigmoidは損失内で）
def forward(self, x):
x = self.features(x)
x = self.pool(x)
x = torch.flatten(x, 1)
x = self.head(x)
return x # ロジット
# ====== 3) 学習・評価ループ ======
@torch.no_grad()
def evaluate(model, loader, device):
model.eval()
total, correct, total_loss = 0, 0, 0.0
criterion = nn.BCEWithLogitsLoss()
for images, targets in loader:
images = images.to(device)
# CIFAR-10のラベルは0/1のままでOK。floatにして BCE 損失に合わせる
targets = targets.float().unsqueeze(1).to(device)
logits = model(images)
loss = criterion(logits, targets)
preds = (torch.sigmoid(logits) >= 0.5).long()
correct += (preds == targets.long()).sum().item()
total += targets.size(0)
total_loss += loss.item() * targets.size(0)
return total_loss / total, correct / total
def train(epochs=5, batch_size=32, lr=1e-3, root="./data", num_workers=2, device=None):
device = device or ("cuda" if torch.cuda.is_available() else "cpu")
train_ds, test_ds = get_cifar10_two_classes(root=root)
train_loader = DataLoader(
train_ds,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
pin_memory=(device == "cuda"),
)
test_loader = DataLoader(
test_ds,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
pin_memory=(device == "cuda"),
)
model = VGG16BinaryClassifier(pretrained=True, freeze_features=True).to(device)
optimizer = torch.optim.Adam(
[p for p in model.parameters() if p.requires_grad], lr=lr
)
criterion = nn.BCEWithLogitsLoss()
hist = {"train_loss": [], "train_acc": [], "val_loss": [], "val_acc": []}
# 外側：エポックのプログレスバー（trangeはtqdm(range)の糖衣）
for epoch in trange(1, epochs + 1, desc="Epochs"):
model.train()
running_loss, running_correct, running_total = 0.0, 0, 0
# 内側：バッチのプログレスバー
with tqdm(
train_loader,
desc=f"Train {epoch}/{epochs}",
unit="batch",
leave=False,
dynamic_ncols=True,
) as pbar:
for images, targets in pbar:
images = images.to(device)
targets = targets.float().unsqueeze(1).to(device)
optimizer.zero_grad()
logits = model(images)
loss = criterion(logits, targets)
loss.backward()
optimizer.step()
preds = (torch.sigmoid(logits) >= 0.5).long()
running_correct += (preds == targets.long()).sum().item()
running_total += targets.size(0)
running_loss += loss.item() * targets.size(0)
# 途中経過をバー右側に表示
pbar.set_postfix(
{
"loss": f"{loss.item():.4f}",
"acc": f"{(running_correct / running_total):.3f}",
}
)
train_loss = running_loss / running_total
train_acc = running_correct / running_total
val_loss, val_acc = evaluate(model, test_loader, device)
hist["train_loss"].append(train_loss)
hist["train_acc"].append(train_acc)
hist["val_loss"].append(val_loss)
hist["val_acc"].append(val_acc)
# エポック終了時の要約はprintでも残す
print(
f"Epoch {epoch}/{epochs} - "
f"loss: {train_loss:.4f} acc: {train_acc:.4f} "
f"val_loss: {val_loss:.4f} val_acc: {val_acc:.4f}"
)
return model, hist
def plot_history(hist):
# Accuracy
plt.figure(figsize=(4, 5))
plt.plot(hist["train_acc"], label="Train")
plt.plot(hist["val_acc"], label="Test")
plt.title("Model accuracy")
plt.ylabel("Accuracy")
plt.xlabel("Epoch")
plt.legend(loc="upper left")
plt.show()
# Loss
plt.figure(figsize=(4, 5))
plt.plot(hist["train_loss"], label="Train")
plt.plot(hist["val_loss"], label="Test")
plt.title("Model loss")
plt.ylabel("Loss")
plt.xlabel("Epoch")
plt.legend(loc="upper left")
plt.show()
if __name__ == "__main__":
# エポック5・バッチ32をデフォルトに
model, hist = train(epochs=5, batch_size=32, lr=1e-3)
plot_history(hist)
1. get_cifar10_two_classes(...)
役割：CIFAR-10からクラス0と1だけを抜き出し、学習・評価用に Subset を返します。
def get_cifar10_two_classes(
root: str = "./data",
classes_keep: List[int] = [0, 1],
) -> Tuple[Subset, Subset]:
"""
CIFAR-10 からクラス0/1のみを抽出した train/test の Subset を返します。
"""
train_transform = transforms.Compose(
[
transforms.Resize((32, 32)), # 念のため固定
transforms.RandomRotation(degrees=90),
transforms.RandomAffine(degrees=0, translate=(0.2, 0.2)),
transforms.RandomHorizontalFlip(p=0.5),
transforms.ToTensor(), # [0,1]スケーリング
# NOTE: ここでNormalizeを入れるとImageNet学習済みと相性◎（後述）
# transforms.Normalize(mean=[0.485,0.456,0.406],
# std=[0.229,0.224,0.225]),
]
)
test_transform = transforms.Compose(
[
transforms.Resize((32, 32)),
transforms.ToTensor(),
]
)
trainset = datasets.CIFAR10(
root=root, train=True, download=True, transform=train_transform
)
testset = datasets.CIFAR10(
root=root, train=False, download=True, transform=test_transform
)
# クラス0/1のみのインデックスを抽出
def filter_indices(ds: datasets.CIFAR10, keep: List[int]) -> List[int]:
# datasets.CIFAR10.targets は list[int]
return [i for i, t in enumerate(ds.targets) if t in keep]
train_idx = filter_indices(trainset, classes_keep)
test_idx = filter_indices(testset, classes_keep)
return Subset(trainset, train_idx), Subset(testset, test_idx)
- 引数
root: データの保存先（初回は自動でDL）。classes_keep: 残すクラスIDのリスト（デフォは[0,1]）。
- 何をしてる？
transformsで 学習時のみ 画像の水増し（回転・平行移動・水平反転）＋ToTensor()で0–1スケーリング。datasets.CIFAR10を読み込んで、targetsからクラスIDフィルタ →Subset化。
- 戻り値：
(train_subset, test_subset) - ここをいじると楽しい
- クラスを増やす →
classes_keep=[0,1,2]など。 - 正規化を入れる →
Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])を train/test 両方に追加。
- クラスを増やす →
2. class VGG16BinaryClassifier(nn.Module)
役割：VGG16の畳み込み部をそのまま特徴抽出器として利用し、その出力に小さなMLP（中間層＋Dropout付き）を載せた2値分類モデル。
class VGG16BinaryClassifier(nn.Module):
def __init__(
self,
pretrained: bool = True,
freeze_features: bool = True,
hidden_dim: int = 128, # ← 追加
dropout: float = 0.5, # ← 追加
):
super().__init__()
base = models.vgg16(
weights=models.VGG16_Weights.IMAGENET1K_V1 if pretrained else None
)
self.features = base.features
if freeze_features:
for p in self.features.parameters():
p.requires_grad = False
self.pool = nn.AdaptiveAvgPool2d((1, 1))
self.head = nn.Sequential(
nn.Linear(512, hidden_dim),
nn.ReLU(inplace=True),
nn.Dropout(p=dropout),
nn.Linear(hidden_dim, 1),
)
def forward(self, x):
x = self.features(x)
x = self.pool(x)
x = torch.flatten(x, 1)
return self.head(x)
構成
self.features
ImageNetで学習済みのVGG16畳み込み部。freeze_features=Trueなら勾配を止め、ヘッドだけ学習。self.poolAdaptiveAvgPool2d((1,1))で特徴マップを1×1に縮約。
CIFAR-10（32×32）入力でも最終的に512次元ベクトルになります。self.headLinear(512 → hidden_dim)ReLUDropout(p=dropout)Linear(hidden_dim → 1)
というMLP構成。出力はロジット1本。
引数
pretrained=True
→ ImageNetで学習済み重みを使う。freeze_features=True
→ 特徴抽出部を固定（ヘッドだけ学習）。hidden_dim=128
→ 中間層のユニット数。複雑さを調整可能。dropout=0.5
→ 過学習防止用のドロップアウト率。
forward の挙動
入力画像 → features → pool → flatten → head → 出力ロジット
- 戻り値は ロジット（未Sigmoid）。
- 損失関数には
BCEWithLogitsLossを使う（内部でSigmoidを計算してくれる）。
注意点
- 入力はRGB（3チャネル）想定なのでCIFAR-10そのままでOK。
hidden_dimやdropoutはデータ量や過学習傾向に応じて調整。- 「線形だけ」のときより表現力が増し、少し複雑な境界も学習しやすくなります。
3. @torch.no_grad() def evaluate(model, loader, device)
役割：検証ループ。損失と精度（正解率）を返す。
@torch.no_grad()
def evaluate(model, loader, device):
model.eval()
total, correct, total_loss = 0, 0, 0.0
criterion = nn.BCEWithLogitsLoss()
for images, targets in loader:
images = images.to(device)
# CIFAR-10のラベルは0/1のままでOK。floatにして BCE 損失に合わせる
targets = targets.float().unsqueeze(1).to(device)
logits = model(images)
loss = criterion(logits, targets)
preds = (torch.sigmoid(logits) >= 0.5).long()
correct += (preds == targets.long()).sum().item()
total += targets.size(0)
total_loss += loss.item() * targets.size(0)
return total_loss / total, correct / total- ポイント
model.eval()で推論モード（Dropout/BNの挙動が評価用に）。@torch.no_grad()で勾配計算をオフ、速くて省メモリ。- 予測は
torch.sigmoid(logits) >= 0.5を閾値に0/1化。
- 戻り値：
(平均損失, 精度)
4. def train(...)
役割：学習の“本丸”。データローダ作成 → モデル構築 → エポック学習＋評価 → ログ蓄積。
def train(epochs=5, batch_size=32, lr=1e-3, root="./data", num_workers=2, device=None):
device = device or ("cuda" if torch.cuda.is_available() else "cpu")
train_ds, test_ds = get_cifar10_two_classes(root=root)
train_loader = DataLoader(
train_ds,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
pin_memory=(device == "cuda"),
)
test_loader = DataLoader(
test_ds,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
pin_memory=(device == "cuda"),
)
model = VGG16BinaryClassifier(pretrained=True, freeze_features=True).to(device)
optimizer = torch.optim.Adam(
[p for p in model.parameters() if p.requires_grad], lr=lr
)
criterion = nn.BCEWithLogitsLoss()
hist = {"train_loss": [], "train_acc": [], "val_loss": [], "val_acc": []}
# 外側：エポックのプログレスバー（trangeはtqdm(range)の糖衣）
for epoch in trange(1, epochs + 1, desc="Epochs"):
model.train()
running_loss, running_correct, running_total = 0.0, 0, 0
# 内側：バッチのプログレスバー
with tqdm(
train_loader,
desc=f"Train {epoch}/{epochs}",
unit="batch",
leave=False,
dynamic_ncols=True,
) as pbar:
for images, targets in pbar:
images = images.to(device)
targets = targets.float().unsqueeze(1).to(device)
optimizer.zero_grad()
logits = model(images)
loss = criterion(logits, targets)
loss.backward()
optimizer.step()
preds = (torch.sigmoid(logits) >= 0.5).long()
running_correct += (preds == targets.long()).sum().item()
running_total += targets.size(0)
running_loss += loss.item() * targets.size(0)
# 途中経過をバー右側に表示
pbar.set_postfix(
{
"loss": f"{loss.item():.4f}",
"acc": f"{(running_correct / running_total):.3f}",
}
)
train_loss = running_loss / running_total
train_acc = running_correct / running_total
val_loss, val_acc = evaluate(model, test_loader, device)
hist["train_loss"].append(train_loss)
hist["train_acc"].append(train_acc)
hist["val_loss"].append(val_loss)
hist["val_acc"].append(val_acc)
# エポック終了時の要約はprintでも残す
print(
f"Epoch {epoch}/{epochs} - "
f"loss: {train_loss:.4f} acc: {train_acc:.4f} "
f"val_loss: {val_loss:.4f} val_acc: {val_acc:.4f}"
)
return model, hist
- 引数
epochs,batch_size,lr: いつもの3点セット。num_workers: DataLoaderの並列度（Windowsなら1–2で安定）。device: 自動判定（CUDAがあればGPU）。
- 流れ
get_cifar10_two_classesでデータ取得。DataLoaderを学習/評価で用意（pin_memoryはGPU時に有効）。VGG16BinaryClassifierを構築し.to(device)。- 最適化対象は
requires_grad=Trueのパラメータ（＝ヘッドのみ）。 - バッチループ：
loss.backward() → optimizer.step()。 evaluateで検証し、histにloss/accを積む。
- 戻り値：
(学習済みmodel, 履歴hist)
5. def plot_history(hist)
役割：学習履歴の可視化（精度と損失のエポック推移）。
def plot_history(hist):
# Accuracy
plt.figure(figsize=(4, 5))
plt.plot(hist["train_acc"], label="Train")
plt.plot(hist["val_acc"], label="Test")
plt.title("Model accuracy")
plt.ylabel("Accuracy")
plt.xlabel("Epoch")
plt.legend(loc="upper left")
plt.show()
# Loss
plt.figure(figsize=(4, 5))
plt.plot(hist["train_loss"], label="Train")
plt.plot(hist["val_loss"], label="Test")
plt.title("Model loss")
plt.ylabel("Loss")
plt.xlabel("Epoch")
plt.legend(loc="upper left")
plt.show()
- 入力：
histは{"train_loss", "train_acc", "val_loss", "val_acc"}を持つ辞書。 - 出力：
matplotlibで2枚の図（Accuracy／Loss）。 - Tips
- 保存したいときは
plt.savefig("acc.png", dpi=150)などを追加。 - Notebookなら
plt.show()はそのままでOK。
- 保存したいときは
6. if __name__ == "__main__":
役割：スクリプト直実行時の入口。
やってること：デフォルトのハイパラで train → plot_history を順に呼ぶだけ。
if __name__ == "__main__":
# エポック5・バッチ32をデフォルトに
model, hist = train(epochs=5, batch_size=32, lr=1e-3)
plot_history(hist)- 実験を変えたいときはここで
epochsやlrをサクッと調整。 - 学習済みモデルの保存は例えば：
torch.save(model.state_dict(), "vgg16_bin_cifar10.pt")
つぶやき：
複数実験を回すなら、argparseでCLI引数にしておくと未来の自分が助かる。
コメント