Pytorch迁移学习识别四种天气

 

In [1]:

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 torchvision.models import VGG16_Weights
from torch.optim import lr_scheduler
import os

 

In [2]:

base_dir = './四种天气图片数据集/'
train_dir = os.path.join(base_dir, 'train')
test_dir = os.path.join(base_dir, 'test')

 

In [3]:

train_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomCrop(192),  # 随机位置剪裁
    transforms.RandomHorizontalFlip(),  # 水平帮转
    transforms.RandomVerticalFlip(),   # 垂直翻转
    transforms.RandomRotation(0.4),  # 随机旋转
    # transforms.ColorJitter(brightness=0.5),   # 亮度
    # transforms.ColorJitter(contrast=0.4),   # 对比度
    # transforms.RandomGrayscale(),   # 随机灰度化
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

test_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

 

In [4]:

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

 

In [5]:

train_ds.class_to_idx

 

Out[5]:

{'cloudy': 0, 'rain': 1, 'shine': 2, 'sunrise': 3}

 

In [6]:

train_dl = torch.utils.data.DataLoader(train_ds, batch_size=32, shuffle=True)
test_dl = torch.utils.data.DataLoader(test_ds, batch_size=64, shuffle=True)

 

In [7]:

imgs, labels = next(iter(train_dl))
imgs.shape

 

Out[7]:

torch.Size([32, 3, 192, 192])

 

In [8]:

labels

 

Out[8]:

tensor([3, 3, 1, 0, 0, 0, 1, 0, 2, 3, 2, 0, 2, 2, 2, 2, 1, 0, 0, 0, 2, 1, 1, 0,
        3, 0, 1, 3, 2, 3, 2, 0])

 

In [9]:

img = imgs[0]
print(img.min(), img.max())
img = (img + 1) / 2
img = img.permute(1, 2, 0)
print(img.min(), img.max())
plt.imshow(img)

 

 

tensor(-0.9922) tensor(1.)
tensor(0.0039) tensor(1.)

Out[9]:

<matplotlib.image.AxesImage at 0x1c26401f760>

 

 

In [10]:

# 加载预训练好的模型
model = torchvision.models.vgg16(weights=VGG16_Weights.DEFAULT)
for param in model.features.parameters():
    param.requires_grad = False

model

 

Out[10]:

VGG(
  (features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace=True)
    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace=True)
    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace=True)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
)

 

In [11]:

# 修改原网络中的输出层结构
model.classifier[-1].out_features = 4
# model.classifier[-1] = torch.nn.Linear(model.classifier[-1], 4)

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

 

In [12]:

def fit(epoch, model, train_loader, test_loader):
    correct = 0
    total = 0
    running_loss = 0
    
    # 因为bn和dropout需要手动指定训练模式和推理模式
    model.train()
    for x, y in train_loader:
        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()

    step_lr_scheduler.step()   # 更新学习率
    epoch_loss = running_loss / len(train_loader.dataset)
    epoch_acc = correct / total
        
    # 测试过程
    test_correct = 0
    test_total = 0
    test_running_loss = 0
    model.eval()
    with torch.no_grad():
        for x, y in test_loader:
            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()
    test_epoch_loss = test_running_loss / len(test_loader.dataset)
    test_epoch_acc = test_correct / test_total

    print('epoch:', epoch, '  loss:', round(epoch_loss, 3), '  accuracy:', round(epoch_acc, 3),
         '  test_loss:', round(test_epoch_loss, 3), '  test_accuracy:', round(test_epoch_acc, 3))
    return epoch_loss, epoch_acc, test_epoch_loss, test_epoch_acc

 

In [13]:

%%time
epochs = 10
train_loss = []
train_acc = []
test_loss = []
test_acc = []
for epoch in range(epochs):
    epoch_loss, epoch_acc, test_epoch_loss, test_epoch_acc = fit(epoch, model, train_dl, test_dl)
    train_loss.append(epoch_loss)
    train_acc.append(epoch_acc)
    test_loss.append(test_epoch_loss)
    test_acc.append(test_epoch_acc)

 

 

epoch: 0   loss: 0.052   accuracy: 0.707   test_loss: 0.005   test_accuracy: 0.916
epoch: 1   loss: 0.013   accuracy: 0.903   test_loss: 0.004   test_accuracy: 0.942
epoch: 2   loss: 0.01   accuracy: 0.93   test_loss: 0.004   test_accuracy: 0.951
epoch: 3   loss: 0.012   accuracy: 0.917   test_loss: 0.011   test_accuracy: 0.84
epoch: 4   loss: 0.012   accuracy: 0.941   test_loss: 0.007   test_accuracy: 0.902
epoch: 5   loss: 0.008   accuracy: 0.961   test_loss: 0.005   test_accuracy: 0.942
epoch: 6   loss: 0.013   accuracy: 0.963   test_loss: 0.004   test_accuracy: 0.951
epoch: 7   loss: 0.006   accuracy: 0.968   test_loss: 0.004   test_accuracy: 0.956
epoch: 8   loss: 0.009   accuracy: 0.968   test_loss: 0.004   test_accuracy: 0.947
epoch: 9   loss: 0.004   accuracy: 0.971   test_loss: 0.002   test_accuracy: 0.96
CPU times: total: 1h 56min 36s
Wall time: 13min 4s

 

In [14]:

plt.plot(range(1, epochs+1), train_loss, label='train_loss')
plt.plot(range(1, epochs+1), test_loss, label='test_loss')
plt.legend()

 

Out[14]:

<matplotlib.legend.Legend at 0x1c26d003f40>