Pytorch实战 | P7咖啡豆图片识别(深度学习实践pytorch)

• 🍨 本文为🔗365天深度学习训练营 中的学习记录博客
• 🍦 参考文章：Pytorch实战 | 第P7周：咖啡豆识别
• 🍖 原作者：K同学啊|接辅导、项目定制

一、我的环境

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

二、主要代码实现

1、main.py

# -*- coding: utf-8 -*-
import copy
import matplotlib.pyplot as plt
import torch.utils.data
from torchvision import datasets, transforms
from vgg16 import VGG16
import torch.nn as nn

# 一、数据准备
# 1、数据导入
train_transforms = transforms.Compose({
    transforms.Resize([224, 224]),
    transforms.RandomHorizontalFlip(),  # 随机水平翻转
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
})

test_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/', transform=train_transforms)
class_to_idx = total_data.class_to_idx
print(f'数据分类为：{class_to_idx}')

# 2、数据集划分
train_size = int(0.8 * len(total_data))
test_size = len(total_data) - train_size
train_dataset, test_dataset = torch.utils.data.random_split(total_data, [train_size, test_size])

# 3、数据分批
batch_size = 32
train_dl = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_dl = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size)

for X, y in test_dl:
    print(&