SEnet实战 -- 垃圾分类

数据集介绍：
有一个文件夹data，下面一个train文件夹，再下面有6个子文件夹，文件夹名称分别是每种垃圾图片的类别，每个类别下面有该类垃圾的图片。
网路结构：
SEnet：resnet18+通道域注意力
采用SEnet网络训练进行分类，加注意力机制后准确率会稍高一些。
SE层结构如下，暂时不讲注意力机制。

训练过程：
训练集占70%，测试集占30%
数据预处理过程比较简单，只进行了大小的调整，全部缩放到224x224
（还可以增加的常规的数据增强操作，如翻转、裁剪等）
&
类别数字代表含义如下：
0：cardboard
1：glass
2：metal
3：paper
4：plastic
5：trash

网络seresnet代码如下：

import torch.nn as nn
import math
import torch.nn.functional as F


def conv3x3(in_planes, out_planes, stride=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=1, bias=False)


class BasicBlock(nn.Module):
    expansion = 1
    
    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)
        self.downsample = downsample
        self.stride = stride
        
        if planes == 64:
            self.globalAvgPool = nn.AvgPool2d(56, stride=1)
        elif planes == 128:
            self.globalAvgPool = nn.AvgPool2d(28, stride=1)
        elif planes == 256:
            self.globalAvgPool = nn.AvgPool2d(14, stride=1)
        elif planes == 512:
            self.globalAvgPool = nn.AvgPool2d(7, stride=1)
        self.fc1 = nn.Linear(in_features=planes, out_features=round(planes / 16))
        self.fc2 = nn.Linear(in_features=round(planes / 16), out_features=planes)
        self.sigmoid = nn.Sigmoid()
    
    def forward(self, x):
        residual = x
        
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        
        out = self.conv2(out)
        out = self.bn2(out)
        
        if self.downsample is not None:
            residual = self.downsample(x)
        
        original_out = out
        out = self.globalAvgPool(out)
        out = out.view(out.size(0), -1)
        out = self.fc1(out)
        out = self.relu(out)
        out = self.fc2(out)
        out = self.sigmoid(out)
        out = out.view(out.size(0), out.size(1), 1, 1)
        out = out * original_out
        
        out += residual
        out = self.relu(out)
        
        return out


class SENet(nn.Module):
    
    def __init__(self, block, layers, num_classes=6):
        self.inplanes = 64
        super(SENet, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
        self.avgpool = nn.AvgPool2d(7, stride=1)
        self.fc = nn.Linear(512 * block.expansion, num_classes)
        
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
    
    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )
        
        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))
        
        return nn.Sequential(*layers)
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        logits = self.fc(x)
        probas = F.softmax(logits, dim=1)
        return logits, probas

训练文件train.py代码如下

import torchvision
import torch
import torch.nn.functional as F
from torchvision import transforms
from wastesorting.seresnet import SENet
from wastesorting.seresnet import BasicBlock
import os
import shutil
import random
import numpy as np
import matplotlib.pyplot as plt

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


class Classification(object):
    def __init__(self, model_name=None, ctx_id=-1):
        self.model_name = model_name
        self.device = torch.device("cuda:" + str(ctx_id)) if ctx_id > -1 else torch.device("cpu")
        self.net = self.load_model()
    
    def load_model(self):
        net = SENet(BasicBlock, [2, 2, 2, 2])
        if self.model_name is not None:
            net.load_state_dict(torch.load(self.model_name, map_location=None if torch.cuda.is_available() else 'cpu'))
        if torch.cuda.is_available():
            net.to(self.device)
        net.eval()
        
        return net
    
    def train(self, dataset=None, batch_size=20, lr=0.05, num_epochs=20):
        train_loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True)
        optimizer = torch.optim.SGD(self.net.parameters(), lr=lr)
        loss_list = []
        train_acc = []
        test_acc = []
        for epoch in range(0, num_epochs):
            self.net.train()
            
            train_loss = 0
            for batch_idx, (features, targets) in enumerate(train_loader):
                features = features.to(device)
                targets = targets.to(device)
                
                logits, probas = self.net.forward(features)
                loss = F.cross_entropy(logits, targets)
                train_loss += loss
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                
                print('Epoch: %03d/%03d | Batch %04d/%04d | Loss: %.4f'
                      % (epoch + 1, num_epochs, batch_idx, len(train_loader), loss))
            
            tr_acc = self.compute_accuracy(train_loader)
            te_acc = self.compute_accuracy(test_loader)
            print('Epoch: %03d/%03d training accuracy: %.2f%% testing accuracy: %.2f%%' % (
                epoch + 1, num_epochs, tr_acc, te_acc))
            
            loss_list.append(train_loss / len(train_loader))
            train_acc.append(tr_acc)
            test_acc.append(te_acc)
            
            if epoch > 15:
                torch.save(self.net.state_dict(), './model' + str(epoch) + '.pth')
        
        l = len(loss_list)
        x = np.arange(0, l)
        plt.figure(figsize=(12, 6))
        ax1 = plt.subplot(121)
        plt.plot(x, loss_list)
        plt.xlabel('epoch')
        plt.ylabel('loss')
        plt.title('loss function curve')
        ax2 = plt.subplot(122)
        plt.plot(x, train_acc, color='r')
        plt.plot(x, test_acc, color='g')
        plt.xlabel('epoch')
        plt.ylabel('accuracy')
        plt.legend(['train_acc', 'test_acc'], loc=4)
        plt.title('accuracy curve')
        plt.savefig("loss_acc.jpg")
    
    def compute_accuracy(self, data_loader):
        correct_pred, num_examples = 0, 0
        for i, (features, targets) in enumerate(data_loader):
            features = features.to(device)
            targets = targets.to(device)
            
            logits, probas = self.net.forward(features)
            _, predicted_labels = torch.max(probas, 1)
            num_examples += targets.size(0)
            correct_pred += (predicted_labels == targets).sum()
        return correct_pred.float() / num_examples * 100
    
    def predict(self, image, transform):
        image_tensor = transform(image).float()
        image_tensor = image_tensor.unsqueeze_(0)
        image_tensor = image_tensor.to(device)
        _, output = self.net(image_tensor)
        _, index = torch.max(output.data, 1)
        return index


if __name__ == '__main__':
    
    # 准备数据
    def train_test_split(img_src_dir, img_to_dir, rate=0.3):
        path_dir = os.listdir(img_src_dir)  # 取图片的原始路径
        file_number = len(path_dir)
        pick_number = int(file_number * rate)  # 按照rate比例从文件夹中取一定数量图片
        sample = random.sample(path_dir, pick_number)  # 随机选取picknumber数量的样本图片
        for name in sample:
            shutil.move(os.path.join(img_src_dir, name), os.path.join(img_to_dir, name))
        return
    
    
    src_dir = './data/train'
    to_dir = './data/test'
    
    if os.path.isdir('./data/test') == False:
        # 添加test文件夹
        os.mkdir('./data/' + 'test')
        for dir in os.listdir(src_dir):
            os.mkdir('./data/test/' + dir)
    
    num = len(os.listdir(os.path.join(to_dir, 'paper')))
    
    if num == 0:
        # 查看图片数量 并分开训练集测试集
        for file in os.listdir(src_dir):
            file_dir = os.path.join(src_dir, file)
            image = os.listdir(file_dir)
            print(file, '图片总量', len(image))
            
            train_test_split(os.path.join(src_dir, file), os.path.join(to_dir, file))
    
    train_dataset = torchvision.datasets.ImageFolder(root=src_dir,
                                                     transform=transforms.Compose([transforms.Resize((224, 224)),
                                                                                   transforms.ToTensor()]))
    test_dataset = torchvision.datasets.ImageFolder(root=to_dir,
                                                    transform=transforms.Compose([transforms.Resize((224, 224)),
                                                                                  transforms.ToTensor()]))
    
    test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=30, shuffle=True)
    
    print(train_dataset.classes)
    print(test_dataset.classes)
    # print(dataset.class_to_idx)
    # print(dataset.imgs)
    
    cls = Classification()
    cls.train(train_dataset)

测试文件test.py代码如下：

from torchvision import transforms
from wastesorting.train import Classification
from PIL import Image

clspre = Classification(model_name='./wastesorting/model18.pth')

transform = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])


# 测试图片
img = Image.open('./wastesorting/data/test/paper/paper20.jpg')

clspre.predict(img, transform)

# {'cardboard': 0, 'glass': 1, 'metal': 2, 'paper': 3, 'plastic': 4, 'trash': 5}