- 🍨 本文为🔗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(&