Pytorch实战 | P6 好莱坞明星图片识别(深度学习实践pytorch)

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

一、我的环境

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

二、主要代码实现

1、main.py

# -*- coding: utf-8 -*-
import copy
import pathlib
from torch import optim
import torch.utils.data
import torchvision.transforms as transforms
from torchvision import datasets
import torch.nn as nn

# 一、前期准备
data_dir = './data/'
data_dir = pathlib.Path(data_dir)
data_paths = list(data_dir.glob('*'))
classeNames = [str(path).split("\\")[1] for path in data_paths]
# 1、将数据处理成dataset
train_transformers = 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_transformers)

# 2、划分数据集
class_to_idx = total_data.class_to_idx

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、将数据处理成dataloader
batch_size = 32
train_dl = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)