导入项目使用的库
========
import torch.optim as optim
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.autograd import Variable
from torchvision.models import densenet121
设置全局参数
======
设置BatchSize、学习率和epochs,判断是否有cuda环境,如果没有设置为cpu
设置全局参数
modellr = 1e-4
BATCH_SIZE = 32
EPOCHS = 5
DEVICE = torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’)
图像预处理
=====
在做图像与处理时,train数据集的transform和验证集的transform分开做,train的图像处理出了resize和归一化之外,还可以设置图像的增强,比如旋转、随机擦除等一系列的操作,验证集则不需要做图像增强,另外不要盲目的做增强,不合理的增强手段很可能会带来负作用,甚至出现Loss不收敛的情况。
数据预处理
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
transform_test = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
读取数据
=====
这个数据集,作者已经做过调整,可以直接使用Python的默认方式读取数据。数据的目录如下图:
训练集中有1.6万张图片,其中有3千多个是秃头,验证集有2万多张,其中秃头是470张,数量差别太大我随机删了一部分。然后,写读取数据的代码。
dataset_train = datasets.ImageFolder(‘Dataset/Train’, transform)
dataset_test = datasets.ImageFolder(‘Dataset/Validation’,transform_test)
读取数据
print(dataset_train.imgs)
导入数据
train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=False)
设置模型