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>
浙公網安備 33010602011771號