Pytorch 深度學習實戰教程（五）：今天，你垃圾分類了嗎？
來自專輯
深度學習教程
點擊上方「Jack Cui」，選擇「加爲星標」
第一時間關注技術乾貨！
git
1
github
垃圾分類
還記得去年，上海如火如荼進行的垃圾分類政策嗎？算法
2020年5月1日起，北京也開始實行「垃圾分類」了！小程序
北京的垃圾分類標準與上海略有差異，垃圾分爲廚餘垃圾、可回收物、有害垃圾和其餘垃圾四大類，分別對應四種不一樣顏色的垃圾桶，即綠色、藍色、紅色和灰色。
微信小程序
繼上海以後，北京也邁入了「垃圾強制分類時代」。微信
垃圾分類，最變態的地方仍是日本。網絡
日本把垃圾分爲資源、可燃、不可燃、危險、塑料、金屬和粗大，這 7 大類垃圾。app
並規定了回收站天天容許回收的垃圾種類，好比如週一收資源的，週二收塑料的。居民要在指定時間、指定地點丟垃圾。像桌子衣櫃這些大件垃圾還要交錢才能扔。less
敢亂扔垃圾的垃圾最多還可能吃 5 年牢飯並罰上 1000 萬日元！dom
不過，有 24 小時在線發牌打理家務的家庭主婦，人家能夠天天花上半小時去搞垃圾分類，而後照樣有時間去刷刷抖音，打打農藥，看看小電影啥的。
如今，中國一線城市的「社畜」們，幹着 996 的活，又要操起「日本主婦的心」。
一家一個碼農就夠慘了，一家雙碼農那就是「慘上加慘」，下班一個比一個晚。
上海實行「垃圾分類」已經快一年了，不知近況如何？
疫情前，曾去上海玩過一次，一個很明顯的感覺是垃圾分類確實是井井有理地進行着，租住的民宿擺放了 4 個垃圾桶，吃剩的垃圾都須要動動腦子才知道怎麼丟。
不過，也在小區附近看到，分類垃圾桶旁，隨意堆放的未分類的垃圾。
北京通知開始實施「垃圾分類」快一個月了，我所居住的小區在市中心，小區內尚未見到分類垃圾箱，卻是公告張貼了不少，應該還處於宣傳階段。
不過公司裏，卻是開始垃圾分類了，看來是從企業開始行動，而後再到我的。
隨着政策的完善，支持力度的加大，不知若干年後，是否會出現一家提供垃圾分類服務的家政公司？
2
垃圾分類助手
吐槽歸吐槽，人們總歸要隨着時代的發展而順勢前行。
好在，一些 APP 或者小程序已經爲咱們準備好了查詢工具。
查詢方式無非三種形式：文字、語音、圖片。
好多家，都有相似的產品。
你能夠查詢 「996加班掉落的頭髮是什麼垃圾」，也能夠查詢「夜宵必備的小龍蝦是什麼垃圾」。
好比，騰訊有個微信小程序，叫「垃圾分類精靈」；
百度 APP 相機識別入口有個 tag ，叫「識垃圾」。
支付寶有個小程序，叫「垃圾分類指南」。
都支持文字、語音、圖片的垃圾分類識別。
上海剛實行垃圾分類的時候，淘寶的「拍立淘」也有垃圾分類識別入口，不過如今貌似已經下線了。
垃圾分類哪家強，體驗一下就知道了。
3
垃圾分類技術
垃圾識別背後的技術是什麼呢？
文字和語音的都相對簡單，文本匹配便可，語音多了一個音頻轉文本的步驟。
基於圖片的垃圾識別就要難很多。
好比，衛生紙你能夠弄成各類形狀，團成一團，或者撕成一條一條。
甚至，能夠把蛋糕惡趣味地作成「便便」的樣子。
讓算法經過圖片去識別這些東西，顯然有些難爲算法。
目前，使用深度學習分類算法去識別垃圾種類，仍是比較難作好的。
通常都是採用多級分類模型或檢索，搭建的超大分類網絡，好比 1 萬多類物體識別，甚至 10 萬。
而後根據類別標籤作映射，映射到最終的垃圾類別。
底層技術實現，其實仍是多分類。
垃圾分類不一樣於通用的圖像識別，通用圖像識別的「魚」，多是一條在水中自由自在嬉戲的金魚。
而垃圾分類識別的「魚」，則極可能是一個躺在餐盤裏僅剩軀幹骨的魚骨頭。
弄個合適的數據集，也是一門技術活。
數據集獲取通常能夠經過如下 3 個渠道：
- 寫爬蟲，爬各大網站的圖片數據，而後使用本身的接口清洗或者人工標註；
- 將需求提交給數據標註團隊，花經費標註數據。
- 前兩個是要麼得有技術、要麼得有錢。
- 最後一個方法，就得碰運氣了。翻論文，找公開數據集，或者去 AI 比賽網站或者 AI 開放平臺碰碰運氣。
比賽，好比能夠去 Kaggle 搜一搜數據集。
AI 開放平臺，能夠去 AI Studio看看。
URL：
https://aistudio.baidu.com/aistudio/datasetoverview
在 AI Studio 我搜索到了不錯的垃圾分類數據集。
一共 56528 張圖片，214 類，總共 7.13 GB。
URL：
https://aistudio.baidu.com/aistudio/datasetdetail/30982
瞧，運氣不錯，找到了一個不錯的數據集。
下載速度也很給力，10 MB/s。
本文使用這個數據集，訓練一個簡單的垃圾分類模型。
4
數據處理
垃圾數據都放在了名字爲「垃圾圖片庫」的文件夾裏。
首先，咱們須要寫個腳本根據文件夾名，生成對應的標籤文件（dir_label.txt）。
前面是小分類標籤，後面是大分類標籤。
而後再將數據集分爲訓練集（train.txt）、驗證集（val.txt）、測試集（test.txt）。
訓練集和驗證集用於訓練模型，測試集用於驗收最終模型效果。
此外，在使用圖片訓練以前還須要檢查下圖片質量，使用 PIL 的 Image 讀取，捕獲 Error 和 Warning 異常，對有問題的圖片直接刪除便可。
寫個腳本生成三個 txt 文件，訓練集 48045 張，驗證集 5652 張，測試集 2826 張。
腳本很簡單，代碼就不貼了，直接提供處理好的文件。
處理好的四個 txt 文件能夠直接下載。
下載地址：
https://github.com/Jack-Cherish/Deep-Learning/tree/master/Pytorch-Seg/lesson-4
將四個 txt 文件放到和「垃圾圖片庫」的相同目錄下便可。
有了前幾篇教程的基礎，寫個數據讀取的代碼應該很輕鬆吧。
編寫 dataset.py 讀取數據，看一下效果。
import torch
from PIL import Image
import os
import glob
from torch.utils.data import Dataset
import random
import torchvision.transforms as transforms
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
class Garbage_Loader(Dataset):
def __init__(self, txt_path, train_flag=True):
self.imgs_info = self.get_images(txt_path)
self.train_flag = train_flag
self.train_tf = transforms.Compose([
transforms.Resize(224),
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
transforms.ToTensor(),
])
self.val_tf = transforms.Compose([
transforms.Resize(224),
transforms.ToTensor(),
])
def get_images(self, txt_path):
with open(txt_path, 'r', encoding='utf-8') as f:
imgs_info = f.readlines()
imgs_info = list(map(lambda x:x.strip().split('\t'), imgs_info))
return imgs_info
def padding_black(self, img):
w, h = img.size
scale = 224. / max(w, h)
img_fg = img.resize([int(x) for x in [w * scale, h * scale]])
size_fg = img_fg.size
size_bg = 224
img_bg = Image.new("RGB", (size_bg, size_bg))
img_bg.paste(img_fg, ((size_bg - size_fg[0]) // 2,
(size_bg - size_fg[1]) // 2))
img = img_bg
return img
def __getitem__(self, index):
img_path, label = self.imgs_info[index]
img = Image.open(img_path)
img = img.convert('RGB')
img = self.padding_black(img)
if self.train_flag:
img = self.train_tf(img)
else:
img = self.val_tf(img)
label = int(label)
return img, label
def __len__(self):
return len(self.imgs_info)
if __name__ == "__main__":
train_dataset = Garbage_Loader("train.txt", True)
print("數據個數：", len(train_dataset))
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=1,
shuffle=True)
for image, label in train_loader:
print(image.shape)
print(label)
讀取 train.txt 文件，加載數據。數據預處理，是將圖片等比例填充到尺寸爲 280 280 的純黑色圖片上，而後再 resize 到 224 224 的尺寸。
這是圖片分類裏，很常規的一種預處理方法。
此外，針對訓練集，使用 pytorch 的 transforms 添加了水平翻轉和垂直翻轉的隨機操做，這也是很常見的一種數據加強方法。
運行結果：
OK，搞定！開始寫訓練代碼！
5
垃圾分類初體驗
咱們使用一個常規的網絡 ResNet50 ，這是一個很是常見的提取特徵的網絡結構。
整個訓練過程也很簡單，訓練步驟不清楚的，能夠看我上兩篇教程：
- Pytorch深度學習實戰教程（三）：UNet模型訓練
- Pytorch深度學習實戰教程（四）：必知必會的煉丹法寶
-
建立 train.py 文件，編寫以下代碼：
from dataset import Garbage_Loader
from torch.utils.data import DataLoader
from torchvision import models
import torch.nn as nn
import torch.optim as optim
import torch
import time
import os
import shutil
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
"""
Author : Jack Cui
Wechat : https://mp.weixin.qq.com/s/OCWwRVDFNslIuKyiCVUoTA
"""
from tensorboardX import SummaryWriter
def accuracy(output, target, topk=(1,)):
"""
計算topk的準確率
"""
with torch.no_grad():
maxk = max(topk)
batch_size = target.size(0)
_, pred = output.topk(maxk, 1, True, True)
pred = pred.t()
correct = pred.eq(target.view(1, -1).expand_as(pred))
class_to = pred[0].cpu().numpy()
res = []
for k in topk:
correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
res.append(correct_k.mul_(100.0 / batch_size))
return res, class_to
def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):
"""
根據 is_best 存模型，通常保存 valid acc 最好的模型
"""
torch.save(state, filename)
if is_best:
shutil.copyfile(filename, 'model_best_' + filename)
def train(train_loader, model, criterion, optimizer, epoch, writer):
"""
訓練代碼
參數：
train_loader - 訓練集的 DataLoader
model - 模型
criterion - 損失函數
optimizer - 優化器
epoch - 進行第幾個 epoch
writer - 用於寫 tensorboardX
"""
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
# switch to train mode
model.train()
end = time.time()
for i, (input, target) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
input = input.cuda()
target = target.cuda()
# compute output
output = model(input)
loss = criterion(output, target)
# measure accuracy and record loss
[prec1, prec5], class_to = accuracy(output, target, topk=(1, 5))
losses.update(loss.item(), input.size(0))
top1.update(prec1[0], input.size(0))
top5.update(prec5[0], input.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % 10 == 0:
print('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
epoch, i, len(train_loader), batch_time=batch_time,
data_time=data_time, loss=losses, top1=top1, top5=top5))
writer.add_scalar('loss/train_loss', losses.val, global_step=epoch)
def validate(val_loader, model, criterion, epoch, writer, phase="VAL"):
"""
驗證代碼
參數：
val_loader - 驗證集的 DataLoader
model - 模型
criterion - 損失函數
epoch - 進行第幾個 epoch
writer - 用於寫 tensorboardX
"""
batch_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
# switch to evaluate mode
model.eval()
with torch.no_grad():
end = time.time()
for i, (input, target) in enumerate(val_loader):
input = input.cuda()
target = target.cuda()
# compute output
output = model(input)
loss = criterion(output, target)
# measure accuracy and record loss
[prec1, prec5], class_to = accuracy(output, target, topk=(1, 5))
losses.update(loss.item(), input.size(0))
top1.update(prec1[0], input.size(0))
top5.update(prec5[0], input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % 10 == 0:
print('Test-{0}: [{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
phase, i, len(val_loader),
batch_time=batch_time,
loss=losses,
top1=top1, top5=top5))
print(' * {} Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}'
.format(phase, top1=top1, top5=top5))
writer.add_scalar('loss/valid_loss', losses.val, global_step=epoch)
return top1.avg, top5.avg
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
if __name__ == "__main__":
# -------------------------------------------- step 1/4 : 加載數據 ---------------------------
train_dir_list = 'train.txt'
valid_dir_list = 'val.txt'
batch_size = 64
epochs = 80
num_classes = 214
train_data = Garbage_Loader(train_dir_list, train_flag=True)
valid_data = Garbage_Loader(valid_dir_list, train_flag=False)
train_loader = DataLoader(dataset=train_data, num_workers=8, pin_memory=True, batch_size=batch_size, shuffle=True)
valid_loader = DataLoader(dataset=valid_data, num_workers=8, pin_memory=True, batch_size=batch_size)
train_data_size = len(train_data)
print('訓練集數量：%d' % train_data_size)
valid_data_size = len(valid_data)
print('驗證集數量：%d' % valid_data_size)
# ------------------------------------ step 2/4 : 定義網絡 ------------------------------------
model = models.resnet50(pretrained=True)
fc_inputs = model.fc.in_features
model.fc = nn.Linear(fc_inputs, num_classes)
model = model.cuda()
# ------------------------------------ step 3/4 : 定義損失函數和優化器等 -------------------------
lr_init = 0.0001
lr_stepsize = 20
weight_decay = 0.001
criterion = nn.CrossEntropyLoss().cuda()
optimizer = optim.Adam(model.parameters(), lr=lr_init, weight_decay=weight_decay)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=lr_stepsize, gamma=0.1)
writer = SummaryWriter('runs/resnet50')
# ------------------------------------ step 4/4 : 訓練 -----------------------------------------
best_prec1 = 0
for epoch in range(epochs):
scheduler.step()
train(train_loader, model, criterion, optimizer, epoch, writer)
# 在驗證集上測試效果
valid_prec1, valid_prec5 = validate(valid_loader, model, criterion, epoch, writer, phase="VAL")
is_best = valid_prec1 > best_prec1
best_prec1 = max(valid_prec1, best_prec1)
save_checkpoint({
'epoch': epoch + 1,
'arch': 'resnet50',
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer' : optimizer.state_dict(),
}, is_best,
filename='checkpoint_resnet50.pth.tar')
writer.close()
代碼並不複雜，網絡結構直接使 torchvision 的 ResNet50 模型，而且採用 ResNet50 的預訓練模型。算法採用交叉熵損失函數，優化器選擇 Adam，並採用 StepLR 進行學習率衰減。
保存模型的策略是選擇在驗證集準確率最高的模型。
batch size 設爲 64，GPU 顯存大約佔 8G，顯存不夠的，能夠調整 batch size 大小。
模型訓練完成，就能夠寫測試代碼了，看下效果吧！
建立 infer.py 文件，編寫以下代碼：
from dataset import Garbage_Loader
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from torchvision import models
import torch.nn as nn
import torch
import os
import numpy as np
import matplotlib.pyplot as plt
#%matplotlib inline
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
def softmax(x):
exp_x = np.exp(x)
softmax_x = exp_x / np.sum(exp_x, 0)
return softmax_x
with open('dir_label.txt', 'r', encoding='utf-8') as f:
labels = f.readlines()
labels = list(map(lambda x:x.strip().split('\t'), labels))
if __name__ == "__main__":
test_list = 'test.txt'
test_data = Garbage_Loader(test_list, train_flag=False)
test_loader = DataLoader(dataset=test_data, num_workers=1, pin_memory=True, batch_size=1)
model = models.resnet50(pretrained=False)
fc_inputs = model.fc.in_features
model.fc = nn.Linear(fc_inputs, 214)
model = model.cuda()
# 加載訓練好的模型
checkpoint = torch.load('model_best_checkpoint_resnet50.pth.tar')
model.load_state_dict(checkpoint['state_dict'])
model.eval()
for i, (image, label) in enumerate(test_loader):
src = image.numpy()
src = src.reshape(3, 224, 224)
src = np.transpose(src, (1, 2, 0))
image = image.cuda()
label = label.cuda()
pred = model(image)
pred = pred.data.cpu().numpy()[0]
score = softmax(pred)
pred_id = np.argmax(score)
plt.imshow(src)
print('預測結果：', labels[pred_id][0])
plt.show()
這裏須要注意的是，DataLoader 讀取的數據須要進行通道轉換，才能顯示。
預測結果：
怎麼樣？還算簡單吧？
趕快訓練一個本身「垃圾分類器」體驗一下吧！
6
總結
- 本文從實戰出發，講解了怎麼訓練一個本身的「垃圾分類器」。
- baseline 已經提供，提高精度，就是一些細節上的優化了。
Jack Cui
推薦搜索
深度學習垃圾分類ResNet50AI
Jack Cui分享或點在看是最大的支持