PyTorchでの転移学習の実装方法は何ですか？

1年 ago
蓮, 翼
1 minute

PyTorchでの転移学習の実装は、通常、以下の手順に従うことができます。

  1. 事前学習済みモデルをロードする：まず、ImageNetなどの大規模データセットで事前に学習されたResNet、VGGなどのモデルをロードします。
  2. 具体なタスクに基づいて事前学習モデルの最後の層またはいくつかの層を変更し、新しいタスクの出力要件に適合させます。
  3. モデルの重みを凍結する：事前学習済みモデルの重みを凍結して、トレーニング中に更新されないようにする。
  4. 新しいタスクに基づいて、新しい損失関数を定義します。
  5. モデルのトレーニング：修正したモデルを新しいデータセットでトレーニングし、追加した層の重みのみを更新します。
  6. モデルの微調整：モデルの性能をさらに向上させる必要がある場合は、一部の事前トレーニングされたモデルの重みを解凍して、モデル全体を継続してトレーニングすることができます。

PyTorchで転移学習を実装する方法を示す簡単なサンプルコードが以下にあります。

import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
import torch.optim as optim
import torch.utils.data as data
from torchvision.datasets import ImageFolder

# 加载预训练模型
pretrained_model = models.resnet18(pretrained=True)

# 修改模型结构
num_ftrs = pretrained_model.fc.in_features
pretrained_model.fc = nn.Linear(num_ftrs, 2)  # 假设新任务是一个二分类问题

# 冻结模型权重
for param in pretrained_model.parameters():
    param.requires_grad = False

# 加载数据
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor()
])
train_dataset = ImageFolder('path_to_train_data', transform=transform)
train_loader = data.DataLoader(train_dataset, batch_size=32, shuffle=True)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(pretrained_model.fc.parameters(), lr=0.001)

# 训练模型
pretrained_model.train()
for epoch in range(10):
    for inputs, labels in train_loader:
        optimizer.zero_grad()
        outputs = pretrained_model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

# 保存模型
torch.save(pretrained_model.state_dict(), 'pretrained_model.pth')

これはシンプルな転移学習の例であり、実際の応用では具体的な状況に応じて調整や最適化が可能です。

bannerAds