Pytorch做图像分类的迁移学习
参考：PyTorch官方教程中文版
实际中，基本没有人会从零开始（随机初始化）训练一个完整的卷积网络，因为相对于网络，很难得到一个足够大的数据集[网络很深, 需要足够大数据集]。通常的做法是在一个很大的数据集上进行预训练得到卷积网络ConvNet, 然后将这个ConvNet的参数作为目标任务的初始化参数或者固定这些参数。
迁移学习的2种场景：
1、微调Convnet：使用预训练的网络（如在imagenet 1000上训练而来的网络）来初始化自己的网络，而不是随机初始化。其他的训练步骤不变。
2、将Convnet看成固定的特征提取器：首先固定ConvNet除了最后的全连接层外的其他所有层。最后的全连接层被替换成一个新的随机 初始化的层，只有这个新的层会被训练[只有这层参数会在反向传播时更新]。
下面代码实现了第二种场景，利用PyTorch进行迁移学习，训练一个模型来对蚂蚁和蜜蜂进行分类。数据集下载：https://download.pytorch.org/tutorial/hymenoptera_data.zip
from __future__ import print_function
from __future__ import division
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data.dataloader import DataLoader
from torchvision import datasets, models, transforms
import time
import os
import copy
if __name__ == "__main__":
# 参数设置
data_dir = "./data/ants_bees" # 文件夹下有2个子文件train和val；每个文件夹下各有2个子文件ants和bees（两类）
model_name = "resnet" # 网络模型名
num_classes = 2 # 类别数
epochs = 15 # 训练epochs
device = 'cpu' # 训练方法：GPU or CPU
input_size = 224 # 输入图片尺寸
# 数据集设置
data_transforms = {
'train': transforms.Compose([
transforms.RandomResizedCrop(input_size), # 随机裁剪然后resize
transforms.RandomHorizontalFlip(), # 随机水平翻转
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # 归一化
]),
'val': transforms.Compose([
transforms.Resize(input_size),
transforms.CenterCrop(input_size),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # 归一化
])}
image_sets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in ['train', 'val']}
data_loaders = {x: DataLoader(image_sets[x], batch_size=8, shuffle=True, num_workers=4) for x in ['train', 'val']}
# 模型设置
model_ft = models.resnet18(pretrained=True) # resnet18模型
for param in model_ft.parameters(): # param: 每一层权重值
param.requires_grad = False # 参数属性设置为不更新
num_feats = model_ft.fc.in_features # 提取的特征维度
model_ft.fc = nn.Linear(num_feats, num_classes) # 修改最后一层fc
model_ft = model_ft.to(device)
params_to_update = []
for name, param in model_ft.named_parameters(): # name:层名; param:权重值
if param.requires_grad:
params_to_update.append(param) # 需要更新的权重参数,即最后一个全连接层:特征维度*类别数
# 优化器和损失函数设置
optimizer = optim.SGD(params_to_update, lr=0.001, momentum=0.9)
criterion = nn.CrossEntropyLoss()
since = time.time()
best_model_wts = copy.deepcopy(model_ft.state_dict()) # 记录最优模型
best_acc = 0.0
for epoch in range(epochs):
print('Epoch {}/{} -----------'.format(epoch, epochs - 1))
for phase in ['train', 'val']:
if phase == 'train':
model_ft.train()
else:
model_ft.eval()
running_loss = 0.0
running_corrects = 0
for inputs, labels in data_loaders[phase]: # 按照batch取数据
inputs = inputs.to(device) # 输入数据
labels = labels.to(device) # 标签, labels = tensor([0, 1, 0, 0, 1, 0, 1, 1])
optimizer.zero_grad() # 零参数梯度
with torch.set_grad_enabled(phase == 'train'):
outputs = model_ft(inputs) # 提取特征
loss = criterion(outputs, labels) # 计算交叉熵损失
_, predict = torch.max(outputs, 1) # 预测label
if phase == 'train':
loss.backward()
optimizer.step()
running_loss += loss.item() * inputs.size(0) # 统计损失函数之和
running_corrects += torch.sum(predict == labels.data) # 计算预测正确样本个数
epoch_loss = running_loss / len(data_loaders[phase].dataset) # 计算所有样本的平均损失函数值
epoch_acc = float(running_corrects) / len(data_loaders[phase].dataset) # 计算所有样本预测的正确率
print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc)) # 打印结果
if phase == 'val' and epoch_acc > best_acc:
best_acc = epoch_acc
best_model_wts = copy.deepcopy(model_ft.state_dict()) # 更新最优模型
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60)) # 训练总时长
print('Best val Acc: {:4f}'.format(best_acc)) # 打印最高测试准确率
model_ft.load_state_dict(best_model_wts) # load最好模型