Pytorch 深度學習實戰教程（五）：今天，你垃圾分類了嗎？

本文 GitHub https://github.com/Jack-Cherish/PythonPark 已收錄，有技術乾貨文章，整理的學習資料，一線大廠面試經驗分享等，歡迎 Star 和 完善。

一、垃圾分類

還記得去年，上海如火如荼進行的垃圾分類政策嗎？

2020年5月1日起，北京也開始實行「垃圾分類」了！

北京的垃圾分類標準與上海略有差別，垃圾分爲廚餘垃圾、可回收物、有害垃圾和其他垃圾四大類，分別對應四種不同顏色的垃圾桶，即綠色、藍色、紅色和灰色。

繼上海之後，北京也邁入了“垃圾強制分類時代”。

垃圾分類，最變態的地方還是日本。

日本把垃圾分爲資源、可燃、不可燃、危險、塑料、金屬和粗大，這 7 大類垃圾。

並規定了回收站每天允許回收的垃圾種類，比如如週一收資源的，週二收塑料的。居民要在指定時間、指定地點丟垃圾。像桌子衣櫃這些大件垃圾還要交錢才能扔。

敢亂扔垃圾的垃圾最多還可能喫 5 年牢飯並罰上 1000 萬日元！

不過，有 24 小時在線發牌打理家務的家庭主婦，人家可以每天花上半小時去搞垃圾分類，然後照樣有時間去刷刷抖音，打打農藥，看看小電影啥的。

現在，中國一線城市的“社畜”們，幹着 996 的活，又要操起“日本主婦的心”。

一家一個碼農就夠慘了，一家雙碼農那就是「慘上加慘」，下班一個比一個晚。

上海實行「垃圾分類」已經快一年了，不知近況如何？

疫情前，曾去上海玩過一次，一個很明顯的感受是垃圾分類確實是井然有序地進行着，租住的民宿擺放了 4 個垃圾桶，喫剩的垃圾都需要動動腦子才知道怎麼丟。

不過，也在小區附近看到，分類垃圾桶旁，隨意堆放的未分類的垃圾。

北京通知開始實施「垃圾分類」有幾個月了，我所居住的小區在市中心，小區內還沒有見到分類垃圾箱，倒是公告張貼了很多，應該還處於宣傳階段。

不過公司裏，倒是開始垃圾分類了，看來是從企業開始行動，然後再到個人。

隨着政策的完善，支持力度的加大，不知若干年後，是否會出現一家提供垃圾分類服務的家政公司？

二、垃圾分類助手

吐槽歸吐槽，人們總歸要隨着時代的發展而順勢前行。

好在，一些 APP 或者小程序已經爲我們準備好了查詢工具。

查詢方式無非三種形式：文字、語音、圖片。

好多家，都有類似的產品。

你可以查詢 「996加班掉落的頭髮是什麼垃圾」，也可以查詢「夜宵必備的小龍蝦是什麼垃圾」。

比如，騰訊有個微信小程序，叫「垃圾分類精靈」；

百度 APP 相機識別入口有個 tag ，叫「識垃圾」。

支付寶有個小程序，叫「垃圾分類指南」。

都支持文字、語音、圖片的垃圾分類識別。

上海剛實行垃圾分類的時候，淘寶的「拍立淘」也有垃圾分類識別入口，不過現在貌似已經下線了。

垃圾分類哪家強，體驗一下就知道了。

三、垃圾分類技術

垃圾識別背後的技術是什麼呢？

文字和語音的都相對簡單，文本匹配即可，語音多了一個音頻轉文本的步驟。

基於圖片的垃圾識別就要難不少。

比如，衛生紙你可以弄成各種形狀，團成一團，或者撕成一條一條。

甚至，可以把蛋糕惡趣味地做成「便便」的樣子。

讓算法通過圖片去識別這些東西，顯然有些難爲算法。

目前，使用深度學習分類算法去識別垃圾種類，還是比較難做好的。

一般都是採用多級分類模型或檢索，搭建的超大分類網絡，比如 1 萬多類物體識別，甚至 10 萬。

然後根據類別標籤做映射，映射到最終的垃圾類別。

底層技術實現，其實還是多分類。

垃圾分類不同於通用的圖像識別，通用圖像識別的「魚」，可能是一條在水中自由自在嬉戲的金魚。

而垃圾分類識別的「魚」，則很可能是一個躺在餐盤裏僅剩軀幹骨的魚骨頭。

弄個合適的數據集，也是一門技術活。

數據集獲取一般可以通過以下 3 個渠道：

• 寫爬蟲，爬各大網站的圖片數據，然後使用自己的接口清洗或者人工標註；
• 將需求提交給數據標註團隊，花經費標註數據。

前兩個是要麼得有技術、要麼得有錢。

• 最後一個方法，就得碰運氣了。翻論文，找公開數據集，或者去 AI 比賽網站或者 AI 開放平臺碰碰運氣。

比賽，比如可以去 Kaggle 搜一搜數據集。

URL：https://www.kaggle.com/

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。

本文使用這個數據集，訓練一個簡單的垃圾分類模型。

四、數據處理

垃圾數據都放在了名字爲「垃圾圖片庫」的文件夾裏。

首先，我們需要寫個腳本根據文件夾名，生成對應的標籤文件（dir_label.txt）。

前面是小分類標籤，後面是大分類標籤。

然後再將數據集分爲訓練集（train.txt）、驗證集（val.txt）、測試集（test.txt）。

訓練集和驗證集用於訓練模型，測試集用於驗收最終模型效果。

此外，在使用圖片訓練之前還需要檢查下圖片質量，使用 PIL 的 Image 讀取，捕獲 Error 和 Warning 異常，對有問題的圖片直接刪除即可。

寫個腳本生成三個 txt 文件，訓練集 48045 張，驗證集 5652 張，測試集 2826 張。

腳本很簡單，代碼就不貼了，直接提供處理好的文件。

處理好的四個 txt 文件可以直接下載。

下載地址：點擊查看

將四個 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，搞定！開始寫訓練代碼！

五、垃圾分類初體驗

我們使用一個常規的網絡 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 讀取的數據需要進行通道轉換，才能顯示。

預測結果：

怎麼樣？還算簡單吧？

趕快訓練一個自己「垃圾分類器」體驗一下吧！

六、總結

• 本文從實戰出發，講解了怎麼訓練一個自己的「垃圾分類器」。
• baseline 已經提供，提升精度，就是一些細節上的優化了。
• 訓練好的模型，關注微信公衆號，後臺回覆「垃圾分類」獲取。

PS：文中出現的所有代碼，均可在我的 github 上下載，歡迎 Follow、Star：點擊查看