【Pytorch】迁移学习——基于别人训练好的模型快速实现自己的分类网络

环境：

Python3.6.8
torch==1.3.1
torchvision==0.4.1

---

完整代码:

import torch
import torch.optim as optim
import torch.nn.functional as F
import torchvision
import torchvision.datasets as datasets
import torchvision.models as models
import torchvision.transforms as transforms
dataset = datasets.ImageFolder(
    'dataset',
    transforms.Compose([
        transforms.ColorJitter(0.1, 0.1, 0.1, 0.1),
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
)
train_dataset, test_dataset = torch.utils.data.random_split(dataset, [len(dataset) - 50, 50])
train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=16,
    shuffle=True,
    num_workers=4
)

test_loader = torch.utils.data.DataLoader(
    test_dataset,
    batch_size=16,
    shuffle=True,
    num_workers=4
)
model = models.alexnet(pretrained=True)
model.classifier[6] = torch.nn.Linear(model.classifier[6].in_features, 2)
device = torch.device('cuda')
model = model.to(device)

NUM_EPOCHS = 30
BEST_MODEL_PATH = 'best_model.pth'
best_accuracy = 0.0

optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

for epoch in range(NUM_EPOCHS):
    for images, labels in iter(train_loader):
        images = images.to(device)
        labels = labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = F.cross_entropy(outputs, labels)
        loss.backward()
        optimizer.step()
    
    test_error_count = 0.0
    for images, labels in iter(test_loader):
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        test_error_count += float(torch.sum(torch.abs(labels - outputs.argmax(1))))
    
    test_accuracy = 1.0 - float(test_error_count) / float(len(test_dataset))
    print('%d: %f' % (epoch, test_accuracy))
    if test_accuracy > best_accuracy:
        torch.save(model.state_dict(), BEST_MODEL_PATH)
        best_accuracy = test_accuracy

代码说明:

import torch 
import torch.optim as optim	#各种优化器
import torch.nn.functional as F	#损失函数
import torchvision
import torchvision.datasets as datasets	#数据集操作
import torchvision.models as models	#包含一些网络模型
import torchvision.transforms as transforms	#图像处理的方法

dataset = datasets.ImageFolder(
    'dataset',	#你的数据集所在文件夹,这里是dataset,根据文件夹中的文件夹进行分类
    transforms.Compose([
        transforms.ColorJitter(0.1, 0.1, 0.1, 0.1),#修改亮度、对比度和饱和度,色调
        transforms.Resize((224, 224)),	#修改尺寸,可比原图大
        transforms.ToTensor(),	#转换为张量
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])	#对数据按通道进行标准化，即先减均值，再除以标准差
    ])
)

# 取所有训练数据中后50个为验证集,其余全部为测试集
train_dataset, test_dataset = torch.utils.data.random_split(dataset, [len(dataset) - 50, 50])

#训练数据加载
train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=16,	#每次读入图片的数量
    shuffle=True,	#改变加载数据的顺序
    num_workers=4	#采用多个线程加载数据,速度更快
)

#测试数据加载
test_loader = torch.utils.data.DataLoader(
    test_dataset,
    batch_size=16,
    shuffle=True,
    num_workers=4
)

model = models.alexnet(pretrained=True)	#使用alexnet模型,并加载它的预训练权重.站在巨人的肩膀上做分类哈哈
model.classifier[6] = torch.nn.Linear(model.classifier[6].in_features, 2)	#修改网络的最后一层,将原本的1000个类修改为我们所需要的类别数量
device = torch.device('cuda')
model = model.to(device)	#使用CUDA进行训练,适用于有GPU的机器

NUM_EPOCHS = 30	
BEST_MODEL_PATH = 'best_model.pth'
best_accuracy = 0.0

optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)	#采用随机梯度下降作为优化器

for epoch in range(NUM_EPOCHS):
    for images, labels in iter(train_loader):
        images = images.to(device)	#网络在哪里数据就在哪里
        labels = labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = F.cross_entropy(outputs, labels)	#计算网络输出和正确标签的Loss,以便进行参数优化
        loss.backward()	#反向传播优化权重
        optimizer.step()
    
    test_error_count = 0.0
    for images, labels in iter(test_loader):
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        test_error_count += float(torch.sum(torch.abs(labels - outputs.argmax(1))))
    
    test_accuracy = 1.0 - float(test_error_count) / float(len(test_dataset))
    print('%d: %f' % (epoch, test_accuracy))
    #只有当测试的准确率得到提高的情况下才会保存新的网络训练权重
    if test_accuracy > best_accuracy:
        torch.save(model.state_dict(), BEST_MODEL_PATH)
        best_accuracy = test_accuracy