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
%matplotlib inline
import torchvision
import os
import shutil
from torchvision import transforms
1
2
3
4
5
6
7
8
9
10
11

数据为1125张图片，放在同一个文件夹中，每一类图片名都是该类

切分数据集

# 创建一个dataste文件夹
base_dir = r'./dataset2'

if not os.path.isdir(base_dir):
    os.mkdir(base_dir)
    # 里面包含训练集和测试集两个文件夹
    train_dir = os.path.join(base_dir, 'train')
    test_dir = os.path.join(base_dir, 'test')
    os.mkdir(train_dir)
    os.mkdir(test_dir)
1
2
3
4
5
6
7
8
9
10

# 在上面创建的文件夹中添加四类图片文件的文件夹
specises = ['cloudy', 'rain', 'shine', 'sunrise']
for train_or_test in ['train', 'test']:
    for spec in specises:
        os.mkdir(os.path.join(base_dir, train_or_test, spec))
1
2
3
4
5

将图片导入到对应的图片文件夹中

image_dir = r'\\四种天气图片数据集\dataset2'
for i, img in enumerate(os.listdir(image_dir)):
    for spec in specises:
        if spec in img:
            s = os.path.join(image_dir, img)
            # 这是按8：2分为训练集和测试集
            if i%5 == 0:
                d = os.path.join(base_dir, 'test', spec, img)
            else:
                d = os.path.join(base_dir, 'train', spec, img)
            shutil.copy(s, d)
1
2
3
4
5
6
7
8
9
10
11

查看切分的数据

for train_or_test in ['train', 'test']:
    for spec in specises:
        print(train_or_test, spec, len(os.listdir(os.path.join(base_dir, train_or_test, spec))))


train cloudy 240
train rain 172
train shine 202
train sunrise 286
test cloudy 60
test rain 43
test shine 51
test sunrise 71
1
2
3
4
5
6
7
8
9
10
11
12
13

数据处理

这是一个数据处理函数，transforms.Compose可以组合几个变换连续一起操作，这里只是将数据转换为tensor

transforms.Normalize( mean = (0.5,0.5,0.5), std = (0.5,0.5,0.5) )并不是指将张量的均值和标准差设为0.5，而是做这么一个运算：输入的每个channel做 ( [0, 1] - mean（0.5) )/ std（0.5）= [-1, 1] 的运算，所以这一句的实际结果是将[0，1]的张量归一化到[-1, 1]上。

transform = transforms.Compose([
                  # 将所有图片转化为96*96的
                  transforms.Resize((96, 96)),
                  # 转换为tensor类型的数据
                  transforms.ToTensor(),
                  # 标准化
                  transforms.Normalize(mean=[0.5, 0.5, 0.5],
                                       std=[0.5, 0.5, 0.5])
])
1
2
3
4
5
6
7
8
9

切分训练集和数据集

ImageFolder 一个通用的数据加载器，数据集中的数据以以下方式组织

root/dog/xx.png
root/dog/xx.png
root/dog/xx.png

root/cat/xx.png
root/cat/xx.png
root/cat/xx.png
1
2
3
4
5
6
7

train_ds = torchvision.datasets.ImageFolder(
               train_dir,
               transform=transform
)
test_ds = torchvision.datasets.ImageFolder(
               test_dir,
               transform=transform
)
1
2
3
4
5
6
7
8

train_ds.classes
1

数据加载器

torch.utils.data.DataLoader数据加载器，结合了数据集和取样器，并且可以提供多个线程处理数据集。

在训练模型时使用到此函数，用来把训练数据分成多个小组，此函数每次抛出一组数据。直至把所有的数据都抛出。就是做一个数据的初始化。

如下第一个函数就是将训练集生成迭代数据，每次迭代的数据为64个，shuffle为洗牌操作，即打乱顺序。

BATCHSIZE = 16
train_dl = torch.utils.data.DataLoader(
                                       train_ds,
                                       batch_size=BATCHSIZE,
                                       shuffle=True
)
test_dl = torch.utils.data.DataLoader(
                                       test_ds,
                                       batch_size=BATCHSIZE,
)
1
2
3
4
5
6
7
8
9
10

train_dl本质上是一个可迭代对象，可以使用iter()进行访问，采用iter(dataloader)返回的是一个迭代器，然后可以使用next()访问。
也可以使用enumerate(dataloader)的形式访问。

imgs, labels = next(iter(train_dl))
imgs.shape,imgs[0].shape

# (torch.Size([16, 3, 96, 96]), torch.Size([3, 96, 96]))
1
2
3
4

# permute就是将数据维度变换一下
im = imgs[0].permute(1, 2, 0)
im.shape

# torch.Size([96, 96, 3])
1
2
3
4
5

数据展示

转换为numpy类型，以方便图像展示

im = im.numpy()
type(im),im.shape

# (numpy.ndarray, (96, 96, 3))
1
2
3
4

转化为字典类型

id_to_class = dict((v, k) for k, v in train_ds.class_to_idx.items())
id_to_class
{0: 'cloudy', 1: 'rain', 2: 'shine', 3: 'sunrise'}
1
2
3

plt.figure(figsize=(12, 8))
for i, (img, label) in enumerate(zip(imgs[:6], labels[:6])):
    img = (img.permute(1, 2, 0).numpy() + 1)/2
    plt.subplot(2, 3, i+1)
    plt.title(id_to_class.get(label.item()))
    plt.imshow(img)
1
2
3
4
5
6

模型搭建

这里就是一些卷积激活池化的操作，最后加了个dropout层

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 3)
        self.conv2 = nn.Conv2d(16, 32, 3)
        self.conv3 = nn.Conv2d(32, 64, 3)
        self.pool = nn.MaxPool2d(2, 2)
        self.drop = nn.Dropout(0.3)
        self.fc1 = nn.Linear(64*10*10, 1024)
        self.fc2 = nn.Linear(1024, 4)
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.pool(x)
        x = F.relu(self.conv2(x))
        x = self.pool(x)
        x = F.relu(self.conv3(x))
        x = self.pool(x)
        x = self.drop(x)
#        print(x.size())
        x = x.view(-1, 64*10*10)
        x = F.relu(self.fc1(x))
        x = self.drop(x)
        x = self.fc2(x)
        return x
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

定义优化函数和损失函数

loss_fn = nn.CrossEntropyLoss()
optim = torch.optim.Adam(model.parameters(), lr=0.001)
1
2

定义训练函数

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)
        optim.zero_grad()
        loss.backward()
        optim.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()
        
    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_accuracy:', round(epoch_test_acc, 3)
             )
        
    return epoch_loss, epoch_acc, epoch_test_loss, epoch_test_acc
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49

开始训练

train_loss = []
train_acc = []
test_loss = []
test_acc = []
epochs = 30
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_loss)
    test_acc.append(epoch_test_acc)
1
2
3
4
5
6
7
8
9
10
11
12
13
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()
1
2
3