Pytorch极简入门教程（十四）——学习速率衰减

学习速率衰减

导入必要的模块

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt

import torchvision
from torchvision import transforms
from torch.utils.data import DataLoader
import os

导入数据

base_dir = r"./dataset/4Weather"
train_dir = os.path.join(base_dir, "train")
test_dir = os.path.join(base_dir, "test")

训练集和测试集的配置

train_transform = transforms.Compose([
    transforms.Resize(224),
    transforms.RandomResizedCrop(192, scale=(0.6, 1.0), ratio=(0.8, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(0.2),
    torchvision.transforms.ColorJitter(brightness=0.5, contrast=0, saturation=0, hue=0),
    torchvision.transforms.ColorJitter(brightness=0.5, contrast=0, saturation=0, hue=0),
    transforms.ToTensor(),
    transforms.Normalize(mean=[.5, .5, .5], std=[0.5,0.5,0.5])
])

train_ds = torchvision.datasets.ImageFolder(
    train_dir,
    transform = train_transform
)

test_transform = transforms.Compose([
    transforms.Resize((192, 192)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5])
])
test_ds = torchvision.datasets.ImageFolder(
    test_dir,
    transform = test_transform
)

BATCH_SIZE = 32
train_dl = DataLoader(
    train_ds,
    batch_size = BATCH_SIZE,
    shuffle = True
)
test_dl = DataLoader(
    test_ds,
    batch_size = BATCH_SIZE,
)

定义VGG16模型

model = torchvision.models.vgg16(pretrained=True, progress=True)
print("model:",model)

冻结基础网络

for para in model.features.parameters():
    para.require_grad = False

model.classifier[-1].out_features =4

定义优化器和学习速度

if torch.cuda.is_available():
    model.to("cuda")
loss_fn = nn.CrossEntropyLoss()
# 学习率衰减
from torch.optim import lr_scheduler
optimizer = torch.optim.Adam(
    model.classifier.parameters(),
    lr = 0.001
)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

定义评判标准

def fit(epoch, model, trainloader, testloader):
    correct = 0
    total = 0
    running_loss = 0
    for x, y in trainloader:
        if torch.cuda.is_available():
            # 获取训练集
            x, y = x.to("cuda"), y.to("cuda")

        # 正向及反向传播
        y_pred = model(x)
        loss = loss_fn(y_pred, y)
        # 权重参数梯度清零
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        with torch.no_grad():
            y_pred = torch.argmax(y_pred, dim=1)
            correct += (y_pred==y).sum().item()
            total += y.size(0)
            running_loss += loss.item()

    exp_lr_scheduler.step()
    epoch_loss = running_loss / len(trainloader.dataset)
    epoch_acc = correct / total

    test_correct = 0
    test_total = 0
    test_running_loss = 0

    with torch.no_grad():
        for x, y in testloader:
            if torch.cuda.is_available():
                x, y = x.to("cuda"), y.to("cuda")
            y_pred = model(x)
            loss = loss_fn(y_pred, y)
            y_pred = torch.argmax(y_pred, dim=1)
            test_correct += (y_pred==y).sum().item()
            test_total += y.size(0) 
            test_running_loss += loss.item()

    epoch_test_loss = test_running_loss / len(testloader.dataset)
    epoch_test_acc = test_correct / test_total

    print("epoch:", epoch,
          "loss:", round(epoch_loss, 3),
          "accuracy:", round(epoch_acc, 3),
          "test_loss:", round(epoch_test_loss, 3),
          "test_acc:", round(epoch_test_acc, 3)
          )
    return epoch_loss, epoch_acc, epoch_test_loss, epoch_test_acc

开启训练和测试

Epochs = 100
train_loss = []
train_acc = []
test_loss =[]
test_acc = []
for epoch in range(Epochs):
    epoch_loss, epoch_acc, epoch_test_loss, epoch_test_acc = fit(epoch, model, train_dl, test_dl)
    train_loss.append(epoch_loss)
    train_acc.append(epoch_acc)
    test_loss.append(epoch_test_acc)
    test_acc.append(epoch_test_acc)

plt.plot(range(1, Epochs+1), train_loss, label="train_loss")
plt.plot(range(1, Epochs+1), test_loss, label="test_loss")
plt.legend()

plt.plot(range(1, Epochs+1), train_acc, label="train_acc")
plt.plot(range(1, Epochs+1), test_acc, label="test_acc")
plt.show()