Pytorch实战 | P8 YOLOv5-C3模块实现天气图片识别

● 🍨 本文为🔗365天深度学习训练营 中的学习记录博客
● 🍦 参考文章：Pytorch实战 | 第P8天：YOLOv5-C3模块实现天气图片识别
● 🍖 原作者：K同学啊|接辅导、项目定制

一、我的环境

● 语言环境：Python3.8
● 编译器：pycharm
● 深度学习环境：Pytorch
● 数据来源：链接: https://pan.baidu.com/s/1SEfd8mvWt7BpzmWOeaIRkQ 提取码: gdie

二、主要代码实现

1、main.py

# -*- coding: utf-8 -*-
import copy

import torch.utils.data
from torchvision import transforms, datasets
from model import model_K, train, test
import torch.nn as nn
import torchsummary as summary

# 1、导入并处理和转换数据

train_transforms = transforms.Compose([
    transforms.Resize([224, 224]),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])
# 加载数据
total_data = datasets.ImageFolder('./data/', train_transforms)
# 划分数据集
train_size = int(len(total_data) * 0.8)
test_size = len(total_data) - train_size
train_data, test_data = torch.utils.data.random_split(total_data, [train_size, test_size])

batch_size = 32
train_dl = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=True)
test_dl = torch.utils.data.DataLoader(test_data, batch_size=batch_size, shuffle=True)

for X, y in test_dl:
    print("Shape of X [N, C, H, W]: ", X.shape)
    print("Shape of y: ", y.shape, y.dtype)
    break

# 2、构建网络模型 model.py

# 3、模型训练
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print("Using {} device".format(device))

model = model_K().to(device)

# 统计模型参数量以及其他指标
summary.summary(model, (3, 224, 224))

optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
loss_fn = nn.CrossEntropyLoss()

epochs = 20

train_l