深度學習實戰之垃圾分類
阿新 • • 發佈:2022-05-11
垃圾分類,指按一定規定或標準將垃圾分類儲存、分類投放和分類搬運,從而轉變成公共資源的一系列活動的總稱。分類的目的是提高垃圾的資源價值和經濟價值,力爭物盡其用;然而我們在日常生活中認為對垃圾分類還是有些不知所措的,對幹垃圾、溼垃圾……分的不是很清楚,由此我們就想到了使用深度學習的方法進行分類。
簡介
本篇博文主要會帶領大家進行資料的預處理、網路搭建、模型訓練、模型測試
1. 獲取資料集
這裡筆者已經為大家提供了一個比較完整的資料集,所以大家不必再自己去收集資料了
資料集連結:https://pan.baidu.com/s/1OhA7yQt9evqNIP5CIPjdgw
提取碼:Z5A1
如下為資料集中的部分資料展示
這裡就不過多的展示了,因為有些圖片比較的不雅
ok,不說題外話了,我們繼續
2.預處理資料集
import torch,visdom,time import os,csv,random,glob from PIL import Image from torchvision import transforms from torch.utils.data import DataLoader,Dataset class Data(Dataset): def __init__(self,root,resize,model): super(Data, self).__init__() self.root=root self.resize=resize # dividing data if model=='train': self.images=self.images[:int(0.6*len(self.images))] self.labels=self.labels[:int(0.6*len(self.labels))] if model=='verify': self.images=self.images[int(0.6*len(self.images)):int(0.8*len(self.images))] self.labels=self.labels[int(0.6*len(self.labels)):int(0.8*len(self.labels))] else: self.images=self.images[int(0.8*len(self.images)):] self.labels=self.labels[int(0.8*len(self.labels)):] def main(): # show data viz = visdom.Visdom() db = Data('data', 64, 'train') x,y = next(iter(db)) print('sample:', x.shape, y.shape, y) viz.image(db.normalize(x), win='sample_x', opts=dict(title='sample_x')) if __name__ == '__main__': main()
這裡需要對資料進行增強,增強的最終結果如下圖所示:
-
搭建網路
接下來進行搭建網路了,筆者這裡使用的是resnet的網路結構
import torch from torch import nn from torch.nn import functional as F class ResBlk(nn.Module): def __init__(self, ch_in, ch_out, stride=1): super(ResBlk, self).__init__() self.conv1 = nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=stride, padding=1) self.bn1 = nn.BatchNorm2d(ch_out) self.conv2 = nn.Conv2d(ch_out, ch_out, kernel_size=3, stride=1, padding=1) self.bn2 = nn.BatchNorm2d(ch_out) self.extra = nn.Sequential() if ch_out != ch_in: # [b, ch_in, h, w] => [b, ch_out, h, w] self.extra = nn.Sequential( nn.Conv2d(ch_in, ch_out, kernel_size=1, stride=stride), nn.BatchNorm2d(ch_out) ) class ResNet18(nn.Module): def __init__(self, num_class): super(ResNet18, self).__init__() self.conv1 = nn.Sequential( nn.Conv2d(3, 16, kernel_size=3, stride=3, padding=0), nn.BatchNorm2d(16) ) # followed 4 blocks # [b, 16, h, w] => [b, 32, h ,w] self.blk1 = ResBlk(16, 32, stride=3) # [b, 32, h, w] => [b, 64, h, w] self.blk2 = ResBlk(32, 64, stride=3) # # [b, 64, h, w] => [b, 128, h, w] self.blk3 = ResBlk(64, 128, stride=2) # # [b, 128, h, w] => [b, 256, h, w] self.blk4 = ResBlk(128, 256, stride=2) self.outlayer = nn.Linear(256*3*3, num_class) def forward(self, x): return x def main(): if __name__ == '__main__': main()
- 訓練
import torch
from torch import optim, nn
import visdom
import torchvision
from torch.utils.data import DataLoader
from Data_Pre import Data
# from resnet import ResNet18
from torchvision.models import resnet18
from utils import Flatten
batchsz = 32
lr = 1e-3
epochs = 20
device = torch.device('cuda')
torch.manual_seed(1234)
train_db = Data('data', 224, model='train')
val_db = Data('data', 224, model='verify')
test_db = Data('data', 224, model='test')
train_loader = DataLoader(train_db, batch_size=batchsz, shuffle=True,
num_workers=4)
val_loader = DataLoader(val_db, batch_size=batchsz, num_workers=2)
test_loader = DataLoader(test_db, batch_size=batchsz, num_workers=2)
viz = visdom.Visdom()
def evalute(model, loader):
model.eval()
correct = 0
total = len(loader.dataset)
for x, y in loader:
x, y = x.to(device), y.to(device)
with torch.no_grad():
logits = model(x)
pred = logits.argmax(dim=1)
correct += torch.eq(pred, y).sum().float().item()
return correct / total
視覺化訓練結果
可以看到經過次的迭代後,模型的效果還是不錯的,哈哈哈哈哈。。。。。。。
5. 測試模型
專案完整程式碼:https://github.com/taotao/garbage
有什麼問題歡迎同筆者討論:1017190168
最後希望大家能夠動手實踐實踐
最新版本:
自筆者於2019年8月發出了深度學習實戰之垃圾分類博文以來,已被訪問上萬次,幾百次的收藏。同時也有很多的人詢問筆者關於這個專案的問題。在這裡非常的感謝大家的厚愛。其實很早就想再迭代一個垃圾分類專案出來,但是由於之前比較的忙,一直沒有時間將其落實。最近沒什麼事情,筆者花了一些時間又重新做了一個垃圾分類系統。希望對大家學習研究有所幫助。
新版本連結:
深度學習之垃圾分類系統(識別率近99%)
垃圾分類識別系統(目標檢測版)
近期專案:
新冠肺炎(Covid-19)病毒檢測
歡迎大家閱讀、獲取。