PyTorch迁移学习-私人数据集上的蚂蚁蜜蜂分类

最新推荐文章于 2024-09-12 22:45:00 发布

Assassin401

最新推荐文章于 2024-09-12 22:45:00 发布

阅读量5k

点赞数 8

分类专栏：学习笔记文章标签： python pytorch 深度学习

本文链接：https://blog.csdn.net/zzh2910/article/details/103987523

版权

本文介绍如何利用PyTorch的迁移学习在私人数据集上进行蚂蚁和蜜蜂的分类。通过微调预训练的resnet18网络，调整其全连接层以适应两类目标，并探讨两种场景：一是微调整个CNN，二是仅训练新的全连接层作为特征提取器。

摘要由CSDN通过智能技术生成

1. 迁移学习的两个主要场景

微调CNN：使用预训练的网络来初始化自己的网络，而不是随机初始化，然后训练即可
将CNN看成固定的特征提取器：固定前面的层，重写最后的全连接层，只有这个新的层会被训练

下面修改预训练好的resnet18网络在私人数据集上进行训练来分类蚂蚁和蜜蜂

2. 数据集下载

这里使用的数据集包含ants和bees训练图片各约120张，验证图片各75张。由于数据样本非常少，如果从0初始化一个网络进行训练很难有令人满意的结果，这时候迁移学习就派上了用场。数据集下载地址，下载后解压到项目目录

3. 导入相关包

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import torchvision
import torchvision.transforms as transforms
import time
import os
import copy

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

4. 加载数据

PyTorch提供了 torchvision.datasets.ImageFolder 方法来加载私人数据集：

# 训练数据集需要扩充和归一化
# 验证数据集仅需要归一化
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