Pytorch官方教程(四)—Transfer Learning Tutorial

最新推荐文章于 2024-02-28 16:47:24 发布

原创

最新推荐文章于 2024-02-28 16:47:24 发布 · 317 阅读

2 ·

CC 4.0 BY-SA版权

2019 April 04 pytorch

Pytorch官方教程(四)—Transfer Learning Tutorial

迁移学习主要的两种应用场景：

微调卷积网络：使用预训练模型初始化网络，而不是随机初始化，比如在imagenet 1000数据集上训练的网络，剩余的正常训练。
卷积网络作为固定的特征提取器：冻结除最后一个全连接层外所有网络的权值。最后一个全连接层被替换为一个具有随机权重的新层，并且只训练这个层。

 
  # License: BSD
# Author: Sasank Chilamkurthy

from __future__ import print_function, division

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import os
import copy

plt.ion()   # interactive mode
 
 

Load Data

使用 torchvision 和 torch.utils.data 加载数据。

训练一个蚂蚁和蜜蜂分类的模型。训练图片每个类有 120 张图片。验证集每个类有 75 张图片。通常，如果从零开始训练，这是一个非常小的数据集。使用转移学习，模型会具有更好的泛化行。

这个数据集是imagenet的一个非常小的子集。

 
  # Data augmentation and normalization for training
# Just normalization for validation
data_transforms = {
     
     
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

data_dir = 'data/hymenoptera_data'
image_datasets = {
     
     x: datasets.ImageFolder(os.path.join(data_dir, x),
                                          data_transforms[x])
                  for x in ['train', 'val']}
dataloaders = {
     
     x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4,
                                             shuffle=True, num_workers=4)
              for x in ['train', 'val']}
dataset_sizes = {

 