深度学习实战之垃圾分类_小言_互联网的博客

深度学习实战之垃圾分类

2021-02-14 09:17 225人阅读评论(0)

垃圾分类，指按一定规定或标准将垃圾分类储存、分类投放和分类搬运，从而转变成公共资源的一系列活动的总称。分类的目的是提高垃圾的资源价值和经济价值，力争物尽其用；然而我们在日常生活中认为对垃圾分类还是有些不知所措的，对干垃圾、湿垃圾……分的不是很清楚，由此我们就想到了使用深度学习的方法进行分类。

简介

本篇博文主要会带领大家进行数据的预处理、网络搭建、模型训练、模型测试

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/huzixuan1/classification
有什么问题欢迎同笔者讨论：1017190168
最后希望大家能够动手实践实践