【圖像分類】實(shí)戰(zhàn)——使用DenseNet實(shí)現(xiàn)識別禿頭（pytorch）

目錄

摘要

導(dǎo)入項(xiàng)目使用的庫

設(shè)置全局參數(shù)

圖像預(yù)處理

讀取數(shù)據(jù)

設(shè)置模型

設(shè)置訓(xùn)練和驗(yàn)證

測試

完整代碼

摘要

我在前面的文章已經(jīng)寫了很多模型的實(shí)戰(zhàn)，這是實(shí)戰(zhàn)的最后一篇了。我沒有加入可視化，也沒有對代碼做過多的裝飾，只希望用最簡單的方式讓大家知道分類模型是怎樣實(shí)現(xiàn)的。

今天我們用DenseNet實(shí)現(xiàn)對禿頭的分類，數(shù)據(jù)集我放在百度網(wǎng)盤了，地址：鏈接：https://pan.baidu.com/s/177ethB_1ZLyl8_Ef1lJxSA 提取碼：47fo 。這個數(shù)據(jù)集可能讓廣大的程序員扎心了。

下面展示一下數(shù)據(jù)集的樣例。

這個都是禿頂?shù)模麄兊墓餐攸c(diǎn)：都是男士，為啥女士不禿頂呢？

導(dǎo)入項(xiàng)目使用的庫

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

設(shè)置全局參數(shù)

設(shè)置BatchSize、學(xué)習(xí)率和epochs，判斷是否有cuda環(huán)境，如果沒有設(shè)置為cpu

# 設(shè)置全局參數(shù)

modellr = 1e-4

BATCH_SIZE = 32

EPOCHS = 5

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

圖像預(yù)處理

在做圖像與處理時，train數(shù)據(jù)集的transform和驗(yàn)證集的transform分開做，train的圖像處理出了resize和歸一化之外，還可以設(shè)置圖像的增強(qiáng)，比如旋轉(zhuǎn)、隨機(jī)擦除等一系列的操作，驗(yàn)證集則不需要做圖像增強(qiáng)，另外不要盲目的做增強(qiáng)，不合理的增強(qiáng)手段很可能會帶來負(fù)作用，甚至出現(xiàn)Loss不收斂的情況。

# 數(shù)據(jù)預(yù)處理

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

])

讀取數(shù)據(jù)

訓(xùn)練集中有1.6萬張圖片，其中有3千多個是禿頭，驗(yàn)證集有2萬多張，其中禿頭是470張，數(shù)量差別太大我隨機(jī)刪了一部分。然后，寫讀取數(shù)據(jù)的代碼。

dataset_train = datasets.ImageFolder('Dataset/Train', transform)

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

# 讀取數(shù)據(jù)

print(dataset_train.imgs)

# 導(dǎo)入數(shù)據(jù)

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)

設(shè)置模型

使用交叉熵作為loss，模型采用densenet121，建議使用預(yù)訓(xùn)練模型，我在調(diào)試的過程中，使用預(yù)訓(xùn)練模型可以快速得到收斂好的模型，使用預(yù)訓(xùn)練模型將pretrained設(shè)置為True即可。更改最后一層的全連接，將類別設(shè)置為2，然后將模型放到DEVICE。優(yōu)化器選用Adam。

# 實(shí)例化模型并且移動到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優(yōu)化器，學(xué)習(xí)率調(diào)低

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

設(shè)置訓(xùn)練和驗(yàn)證

最外層是循環(huán)的是每個epochs，先訓(xùn)練，后驗(yàn)證。下面分別講一下訓(xùn)練和驗(yàn)證的過程。

訓(xùn)練過程必須經(jīng)歷的步驟：

第一步：將輸入input向前傳播，進(jìn)行運(yùn)算后得到輸出output，代碼：output = model(data)

第二步：將output再輸入loss函數(shù)，計(jì)算loss值（是個標(biāo)量），代碼：? loss = criterion(output, target)

第三步：將梯度反向傳播到每個參數(shù)，代碼：? loss.backward()

第四步：將參數(shù)的grad值初始化為0，代碼：?optimizer.zero_grad()

第五步：更新權(quán)重，代碼：?optimizer.step()

驗(yàn)證過程和訓(xùn)練過程基本相似。

# 定義訓(xùn)練過程

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

# 驗(yàn)證過程

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

# 訓(xùn)練

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了，運(yùn)行結(jié)果如下：

測試

測試集存放的目錄如下圖：

第一步 定義類別，這個類別的順序和訓(xùn)練時的類別順序?qū)?yīng)，一定不要改變順序！！！！我們在訓(xùn)練時，Bald類別是0，NoBald類別是1，所以我定義classes為('Bald','NoBald')。

第二步 定義transforms，transforms和驗(yàn)證集的transforms一樣即可，別做數(shù)據(jù)增強(qiáng)。

第三步 加載model，并將模型放在DEVICE里，

第四步 讀取圖片并預(yù)測圖片的類別，在這里注意，讀取圖片用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()]))

運(yùn)行結(jié)果如下：

第二種方法可以使用pytorch默認(rèn)加載數(shù)據(jù)集的方法。

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

# 對應(yīng)文件夾的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

運(yùn)行結(jié)果：

完整代碼

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

# 設(shè)置全局參數(shù)

modellr = 1e-4

BATCH_SIZE = 32

EPOCHS = 5

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

# 數(shù)據(jù)預(yù)處理

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)

# 讀取數(shù)據(jù)

print(dataset_train.imgs)

# 導(dǎo)入數(shù)據(jù)

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)

# 實(shí)例化模型并且移動到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優(yōu)化器，學(xué)習(xí)率調(diào)低

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

# 定義訓(xùn)練過程

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

# 驗(yàn)證過程

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

# 訓(xùn)練

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-深度學(xué)習(xí)文檔類資源-CSDN下載

pytorch 機(jī)器學(xué)習(xí)

版權(quán)聲明：本文內(nèi)容由網(wǎng)絡(luò)用戶投稿，版權(quán)歸原作者所有，本站不擁有其著作權(quán)，亦不承擔(dān)相應(yīng)法律責(zé)任。如果您發(fā)現(xiàn)本站中有涉嫌抄襲或描述失實(shí)的內(nèi)容，請聯(lián)系我們jiasou666@gmail.com 處理，核實(shí)后本網(wǎng)站將在24小時內(nèi)刪除侵權(quán)內(nèi)容。