【图像分类】实战——使用DenseNet实现识别秃头（pytorch）

目录

摘要

导入项目使用的库

设置全局参数

图像预处理

读取数据

设置模型

设置训练和验证

测试

完整代码

摘要

我在前面的文章已经写了很多模型的实战，这是实战的最后一篇了。我没有加入可视化，也没有对代码做过多的装饰，只希望用最简单的方式让大家知道分类模型是怎样实现的。

今天我们用DenseNet实现对秃头的分类，数据集我放在百度网盘了，地址：链接：https://pan.baidu.com/s/177ethB_1ZLyl8_Ef1lJxSA 提取码：47fo 。这个数据集可能让广大的程序员扎心了。

下面展示一下数据集的样例。

这个都是秃顶的，他们的共同特点：都是男士，为啥女士不秃顶呢？

导入项目使用的库

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])

])

读取数据

训练集中有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)

设置模型

使用交叉熵作为loss，模型采用densenet121，建议使用预训练模型，我在调试的过程中，使用预训练模型可以快速得到收敛好的模型，使用预训练模型将pretrained设置为True即可。更改最后一层的全连接，将类别设置为2，然后将模型放到DEVICE。优化器选用Adam。

# 实例化模型并且移动到GPU

criterion = nn.CrossEntropyLoss()

model_ft = densenet121(pretrained=True)

num_ftrs = model_ft.classifier.in_features

model_ft.classifier = nn.Linear(num_ftrs, 2)

model_ft.to(DEVICE)

# 选择简单暴力的Adam优化器，学习率调低

optimizer = optim.Adam(model_ft.parameters(), lr=modellr)

def adjust_learning_rate(optimizer, epoch):

"""Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""

modellrnew = modellr * (0.1 ** (epoch // 50))

print("lr:", modellrnew)

for param_group in optimizer.param_groups:

param_group['lr'] = modellrnew

设置训练和验证

最外层是循环的是每个epochs，先训练，后验证。下面分别讲一下训练和验证的过程。

训练过程必须经历的步骤：

第一步：将输入input向前传播，进行运算后得到输出output，代码：output = model(data)

第二步：将output再输入loss函数，计算loss值（是个标量），代码：  loss = criterion(output, target)

第三步：将梯度反向传播到每个参数，代码：  loss.backward()

第四步：将参数的grad值初始化为0，代码： optimizer.zero_grad()

第五步：更新权重，代码： optimizer.step()

验证过程和训练过程基本相似。

# 定义训练过程

def train(model, device, train_loader, optimizer, epoch):

model.train()

sum_loss = 0

total_num = len(train_loader.dataset)

print(total_num, len(train_loader))

for batch_idx, (data, target) in enumerate(train_loader):

data, target = Variable(data).to(device), Variable(target).to(device)

output = model(data)

loss = criterion(output, target)

optimizer.zero_grad()

loss.backward()

optimizer.step()

print_loss = loss.data.item()

sum_loss += print_loss

if (batch_idx + 1) % 50 == 0:

print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(

epoch, (batch_idx + 1) * len(data), len(train_loader.dataset),

100. * (batch_idx + 1) / len(train_loader), loss.item()))

ave_loss = sum_loss / len(train_loader)

print('epoch:{},loss:{}'.format(epoch, ave_loss))

# 验证过程

def val(model, device, test_loader):

model.eval()

test_loss = 0

correct = 0

total_num = len(test_loader.dataset)

print(total_num, len(test_loader))

with torch.no_grad():

for data, target in test_loader:

data, target = Variable(data).to(device), Variable(target).to(device)

output = model(data)

loss = criterion(output, target)

_, pred = torch.max(output.data, 1)

correct += torch.sum(pred == target)

print_loss = loss.data.item()

test_loss += print_loss

correct = correct.data.item()

acc = correct / total_num

avgloss = test_loss / len(test_loader)

print('\nVal set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(

avgloss, correct, len(test_loader.dataset), 100 * acc))

# 训练

for epoch in range(1, EPOCHS + 1):

adjust_learning_rate(optimizer, epoch)

train(model_ft, DEVICE, train_loader, optimizer, epoch)

val(model_ft, DEVICE, test_loader)

torch.save(model_ft, 'model.pth')

完成后就可以run了，运行结果如下：

测试

测试集存放的目录如下图：

第一步 定义类别，这个类别的顺序和训练时的类别顺序对应，一定不要改变顺序！！！！我们在训练时，Bald类别是0，NoBald类别是1，所以我定义classes为('Bald','NoBald')。

第二步 定义transforms，transforms和验证集的transforms一样即可，别做数据增强。

第三步 加载model，并将模型放在DEVICE里，

第四步 读取图片并预测图片的类别，在这里注意，读取图片用PIL库的Image。不要用cv2，transforms不支持。

import torch.utils.data.distributed

import torchvision.transforms as transforms

from PIL import Image

from torch.autograd import Variable

import os

classes=('Bald','NoBald')

transform_test = transforms.Compose([

transforms.Resize((224, 224)),

transforms.ToTensor(),

transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])

])

DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model = torch.load("model.pth")

model.eval()

model.to(DEVICE)

path='Dataset/Test/Bald/'

testList=os.listdir(path)

for file in testList:

img=Image.open(path+file)

img=transform_test(img)

img.unsqueeze_(0)

img = Variable(img).to(DEVICE)

out=model(img)

# Predict

_, pred = torch.max(out.data, 1)

print('Image Name:{},predict:{}'.format(file,classes[pred.data.item()]))

运行结果如下：

第二种方法可以使用pytorch默认加载数据集的方法。

import torch.utils.data.distributed

import torchvision.transforms as transforms

import torchvision.datasets as datasets

from torch.autograd import Variable

classes=('Bald','NoBald')

transform_test = transforms.Compose([

transforms.Resize((224, 224)),

transforms.ToTensor(),

transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])

])

DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model = torch.load("model.pth")

model.eval()

model.to(DEVICE)

dataset_test = datasets.ImageFolder('Dataset/Test',transform_test)

print(len(dataset_test))

# 对应文件夹的label

for index in range(len(dataset_test)):

item = dataset_test[index]

img, label = item

img.unsqueeze_(0)

data = Variable(img).to(DEVICE)

output = model(data)

_, pred = torch.max(output.data, 1)

print('Image Name:{},predict:{}'.format(dataset_test.imgs[index], classes[pred.data.item()]))

index += 1

运行结果：

完整代码

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

# 设置全局参数

modellr = 1e-4

BATCH_SIZE = 32

EPOCHS = 5

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 数据预处理

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])

])

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)

# 实例化模型并且移动到GPU

criterion = nn.CrossEntropyLoss()

model_ft = densenet121(pretrained=True)

num_ftrs = model_ft.classifier.in_features

model_ft.classifier = nn.Linear(num_ftrs, 2)

model_ft.to(DEVICE)

# 选择简单暴力的Adam优化器，学习率调低

optimizer = optim.Adam(model_ft.parameters(), lr=modellr)

def adjust_learning_rate(optimizer, epoch):

"""Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""

modellrnew = modellr * (0.1 ** (epoch // 50))

print("lr:", modellrnew)

for param_group in optimizer.param_groups:

param_group['lr'] = modellrnew

# 定义训练过程

def train(model, device, train_loader, optimizer, epoch):

model.train()

sum_loss = 0

total_num = len(train_loader.dataset)

print(total_num, len(train_loader))

for batch_idx, (data, target) in enumerate(train_loader):

data, target = Variable(data).to(device), Variable(target).to(device)

output = model(data)

loss = criterion(output, target)

optimizer.zero_grad()

loss.backward()

optimizer.step()

print_loss = loss.data.item()

sum_loss += print_loss

if (batch_idx + 1) % 50 == 0:

print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(

epoch, (batch_idx + 1) * len(data), len(train_loader.dataset),

100. * (batch_idx + 1) / len(train_loader), loss.item()))

ave_loss = sum_loss / len(train_loader)

print('epoch:{},loss:{}'.format(epoch, ave_loss))

# 验证过程

def val(model, device, test_loader):

model.eval()

test_loss = 0

correct = 0

total_num = len(test_loader.dataset)

print(total_num, len(test_loader))

with torch.no_grad():

for data, target in test_loader:

data, target = Variable(data).to(device), Variable(target).to(device)

output = model(data)

loss = criterion(output, target)

_, pred = torch.max(output.data, 1)

correct += torch.sum(pred == target)

print_loss = loss.data.item()

test_loss += print_loss

correct = correct.data.item()

acc = correct / total_num

avgloss = test_loss / len(test_loader)

print('\nVal set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(

avgloss, correct, len(test_loader.dataset), 100 * acc))

# 训练

for epoch in range(1, EPOCHS + 1):

adjust_learning_rate(optimizer, epoch)

train(model_ft, DEVICE, train_loader, optimizer, epoch)

val(model_ft, DEVICE, test_loader)

torch.save(model_ft, 'model.pth')

DenseNet图像分类.zip-深度学习文档类资源-CSDN下载

pytorch 机器学习

版权声明：本文内容由网络用户投稿，版权归原作者所有，本站不拥有其著作权，亦不承担相应法律责任。如果您发现本站中有涉嫌抄袭或描述失实的内容，请联系我们jiasou666@gmail.com 处理，核实后本网站将在24小时内删除侵权内容。