Pytorch---使用Pytorch实现四种天气分类问题

一、代码中的数据集可以通过以下链接获取

百度网盘提取码：lala

二、代码运行环境

Pytorch-gpu==1.7.1
Python==3.7

三、数据集处理代码如下所示

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


def loader_data():
    BATCH_SIZE = 64
    transform = transforms.Compose([
        transforms.Resize((96, 96)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5],
                             std=[0.5, 0.5, 0.5])
    ])
    train_ds = torchvision.datasets.ImageFolder(root=os.path.join('dataset', 'train_weather'), transform=transform)
    test_ds = torchvision.datasets.ImageFolder(root=os.path.join('dataset', 'test_weather'), transform=transform)
    train_dl = DataLoader(dataset=train_ds, batch_size=BATCH_SIZE, shuffle=True)
    test_dl = DataLoader(dataset=test_ds, batch_size=BATCH_SIZE)

    return train_dl, test_dl, test_ds.class_to_idx

四、模型的构建代码如下所示

import torch
from torch import nn


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3)
        self.bn1 = nn.BatchNorm2d(num_features=16)
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(num_features=32)
        self.conv3 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn3 = nn.BatchNorm2d(num_features=64)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.drop = nn.Dropout(0.5)
        self.drop2d = nn.Dropout2d(0.5)
        self.fc1 = nn.Linear(in_features=64 * 10 * 10, out_features=1024)
        self.bn_f1 = nn.BatchNorm1d(num_features=1024)
        self.fc2 = nn.Linear(in_features=1024, out_features=256)
        self.bn_f2 = nn.BatchNorm1d(num_features=256)
        self.fc3 = nn.Linear(in_features=256, out_features=4)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.bn1(x)
        x = self.pool(torch.relu(self.conv2(x)))
        x = self.bn2(x)
        x = self.pool(torch.relu(self.conv3(x)))
        x = self.bn3(x)
        x = self.drop2d(x)
        x = x.view(-1, x.size(1) * x.size(2) * x.size(3))
        x = torch.relu(self.fc1(x))
        x = self.bn_f1(x)
        x = self.drop(x)
        x = torch.relu(self.fc2(x))
        x = self.bn_f2(x)
        x = self.drop(x)
        x = self.fc3(x)
        return x

五、模型的训练代码如下所示

import torch
from data_loader import loader_data
from model_loader import Net
from torch import nn
import tqdm
from sklearn.metrics import accuracy_score
import numpy as np
import os

# 数据的加载
train_dl, test_dl, class_index = loader_data()

# 模型的加载
model = Net()

# 训练的配置
loss_fn = nn.CrossEntropyLoss()
optim = torch.optim.Adam(params=model.parameters(), lr=0.001)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
EPOCHS = 100

# 进行训练
model = model.to(device)
for epoch in range(EPOCHS):
    # 训练部分
    model.train()
    train_tqdm = tqdm.tqdm(iterable=train_dl, total=len(train_dl))
    train_tqdm.set_description_str('Train epoch {:2d}'.format(epoch))
    train_accuracy_sum = []
    train_loss_sum = []
    for images, labels in train_tqdm:
        images, labels = images.to(device), labels.to(device)
        pred = model(images)
        loss = loss_fn(pred, labels)
        optim.zero_grad()
        loss.backward()
        optim.step()
        # 进行训练结果的展示
        train_loss_sum.append(loss.item())
        pred = torch.argmax(input=pred, dim=-1)
        train_accuracy_sum.append(accuracy_score(y_true=labels.cpu().numpy(), y_pred=pred.cpu().numpy()))
        train_tqdm.set_postfix_str(
            'loss is {:14f}, accuracy is {:14f}'.format(np.mean(train_loss_sum), np.mean(train_accuracy_sum)))
    train_tqdm.close()

    # 验证部分
    with torch.no_grad():
        model.eval()
        test_tqdm = tqdm.tqdm(iterable=test_dl, total=len(test_dl))
        test_tqdm.set_description_str('Val epoch {:2d}'.format(epoch))
        test_accuracy_sum = []
        test_loss_sum = []
        for images, labels in test_tqdm:
            images, labels = images.to(device), labels.to(device)
            pred = model(images)
            loss = loss_fn(pred, labels)
            # 进行训练结果的展示
            test_loss_sum.append(loss.item())
            pred = torch.argmax(input=pred, dim=-1)
            test_accuracy_sum.append(accuracy_score(y_true=labels.cpu().numpy(), y_pred=pred.cpu().numpy()))
            test_tqdm.set_postfix_str(
                'loss is {:14f}, accuracy is {:14f}'.format(np.mean(test_loss_sum), np.mean(test_accuracy_sum)))
        test_tqdm.close()

# 模型的保存
if not os.path.exists(os.path.join('model_data')):
    os.mkdir(os.path.join('model_data'))
torch.save(model.state_dict(), os.path.join('model_data', 'model.pth'))

六、模型的预测代码如下所示

import os
import torch
from data_loader import loader_data
from model_loader import Net
import matplotlib.pyplot as plt
import matplotlib

# 数据的加载
train_dl, test_dl, class_index = loader_data()
image, label = next(iter(test_dl))
new_class = dict((v, k) for k, v in class_index.items())

# 模型的加载
model = Net()
model_state_dict = torch.load(os.path.join('model_data', 'model.pth'))
model.load_state_dict(model_state_dict)
model.eval()

# 进行模型的预测
index = 23
with torch.no_grad():
    pred = model(image)
    pred = torch.argmax(input=pred, dim=-1)
    matplotlib.rc("font", family='Microsoft YaHei')
    plt.axis('off')
    plt.title('预测结果为：' + new_class.get(pred[index].item()) + '，标签结果为：' + new_class.get(label[index].item()))
    plt.imshow(image[index].permute(1, 2, 0))
    plt.savefig('result.png')
    plt.show()

七、代码的运行结果如下所示