SEnet实战 -- 垃圾分类_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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

训练文件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)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155

测试文件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}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15