阅读 26

【技术博客】基于AlexNet网络的垃圾分类

【技术博客】基于AlexNet网络的垃圾分类

AlexNet

AlexNet模型来源于论文-ImageNet Classification with Deep Convolutional Neural Networks,作者Alex Krizhevsky,Ilya Sutskever,Geoffrey E.Hinton. AlexNet在ImageNet LSVRC-2012比赛中,达到最低的15.3%的Top-5错误率,比第二名低10.8个百分点。

网络结构

AlexNet包含八层,前五层是卷积层,最后三层是全连接层。它使用了ReLU激活函数,显示出比tabh和sigmoid更好的训练性能。

论文中的图比较抽象,不便于分析结构,下面提供一个更直观的结构图。

参考链接:Netscope

  • 第一层(卷积层)

输入数据:227×227×3 卷积核:11×11×3;步长:4;数量:96 卷积后数据:55×55×96 relu后的数据:55×55×96 Max pool的核:3×3,步长:2 Max pool后的数据:27×27×96 norm1:local_size=5 (LRN(Local Response Normalization) 局部响应归一化) 最后的输出:27×27×96

  • 第二层(卷积层)

输入数据:27×27×96 卷积核:5×5;步长:1;数量:256 卷积后数据:27×27×256 (做了Same padding(相同补白),使得卷积后图像大小不变。) relu2后的数据:27×27×256 Max pool2的核:3×3,步长:2 Max pool2后的数据:13×13×256 ((27-3)/2+1=13 ) norm2:local_size=5 (LRN(Local Response Normalization) 局部响应归一化) 最后的输出:13×13×256

  • 第三层(卷积层)

输入数据:13×13×256 卷积核:3×3;步长:1;数量(也就是输出个数):384 卷积后数据:13×13×384 (做了Same padding(相同补白),使得卷积后图像大小不变。) relu3后的数据:13×13×384 最后的输出:13×13×384 第三层没有Max pool层和norm层

  • 第四层(卷积层)

输入数据:13×13×384 卷积核:3×3;步长:1;数量(也就是输出个数):384 卷积后数据:13×13×384 (做了Same padding(相同补白),使得卷积后图像大小不变。) relu4后的数据:13×13×384 最后的输出:13×13×384 第四层没有Max pool层和norm层

  • 第五层(卷积层)

输入数据:13×13×384 卷积核:3×3;步长:1;数量(也就是输出个数):256 卷积后数据:13×13×256 (做了Same padding(相同补白),使得卷积后图像大小不变。) relu5后的数据:13×13×256 Max pool5的核:3×3,步长:2 Max pool2后的数据:6×6×256 ((13-3)/2+1=6 ) 最后的输出:6×6×256 第五层有Max pool,没有norm层

  • 第六层(全连接层)

输入数据:6×6×256 全连接输出:4096×1 relu6后的数据:4096×1 drop out6后数据:4096×1 最后的输出:4096×1

  • 第七层(全连接层)

输入数据:4096×1 全连接输出:4096×1 relu7后的数据:4096×1 drop out7后数据:4096×1 最后的输出:4096×1

  • 第八层(全连接层)

输入数据:4096×1 全连接输出:1000 fc8输出一千种分类的概率。

数据集预处理

本次实验中使用的垃圾分类数据集一共2307张图片,分为六个分类, cardboard(370), glass(457), metal(380), paper(540), plastic(445), trash(115)。数据集中的图片是经过处理的512x384的三通道图片。 由于该数据集较小,因此需要通过数据增强扩充数据集。在本次试验中通过对图片进行随机翻转,裁剪227x227大小的子图扩充数据集,为了提高模型的准确率,在输入模型前,还需要对图片进行归一化处理,将每个像素的值映射到(0,1)之间。 定义一个torch.utils.data.Dataset类的子类,用于从硬盘中加载数据集,因为存在随机裁剪,在GarbageDataset类中将数据集大小扩大10倍。

class GarbageDataset(Dataset):

    classifications = ["cardboard", "glass", "metal", "paper", "plastic", "trash"]

    def __init__(self, root_dir, transform = None):
        super(GarbageDataset, self).__init__()
        self.root_dir = root_dir
        self.transform = transform
        self.imgs = []
        self.read()

    def __len__(self):
        return 10 * len(self.imgs)

    def __getitem__(self, item):
        img, label = self.imgs[item % len(self.imgs)]
        if self.transform:
            img = self.transform(img)
        return img, label

    def read(self):
        img_dir = os.path.join(self.root_dir, "garbage")
        for i, c in enumerate(GarbageDataset.classifications, 0):
            dir = os.path.join(img_dir, c)
            for img_name in os.listdir(dir):
                img = Image.open(os.path.join(dir, img_name))
                self.imgs.append((img, i))
复制代码

定义transforms, 实例化GarbageDataset加载数据集, 并按照6:2:2的比例划分训练集,验证集和测试集。

dataset = GarbageDataset("data", transform=transforms.Compose([
    transforms.Resize(227),
    transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(227),
    transforms.RandomRotation(90),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
]))
dataset_size = len(dataset)
validset_size = int(dataset_size / 5)
testset_size = validset_size
trainset_size = dataset_size - validset_size - testset_size

trainset, validset, testset = torch.utils.data.random_split(dataset, [trainset_size, validset_size, testset_size])
复制代码

对训练集,验证集,测试集分别实例化一个DataLoader

# 训练集需要打乱顺序
trainloader = DataLoader(dataset=trainset, batch_size=128, shuffle=True)
# 验证集和测试集可以不用打乱数据顺序
validloader = DataLoader(dataset=validset, batch_size=128, shuffle=False)
testloader = DataLoader(dataset=testset, batch_size=128, shuffle=False)
复制代码

模型搭建

定义模型

class GarbageNet(nn.Module):

    def __init__(self):
        super(GarbageNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 96, 11, 4)
        self.conv2 = nn.Conv2d(96, 256, 5, 1, padding=2, groups=2)
        self.conv3 = nn.Conv2d(256, 384, 3, 1, padding=1)
        self.conv4 = nn.Conv2d(384, 384, 3, 1, padding=1)
        self.conv5 = nn.Conv2d(384, 256, 3, 1, padding=1)
        self.fc1 = nn.Linear(256 * 6 * 6, 4096)
        self.fc2 = nn.Linear(4096, 4096)
        self.fc3 = nn.Linear(4096, 6)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = F.max_pool2d(x, kernel_size=3, stride=2)
        # x = F.max_pool2d(F.relu(self.conv1(x)), kernel_size=3, stride=2)
        x = F.max_pool2d(F.relu(self.conv2(x)), kernel_size=3, stride=2)
        x = F.relu(self.conv3(x))
        x = F.relu(self.conv4(x))
        x = F.max_pool2d(self.conv5(x), kernel_size=3, stride=2)
        x = x.view(-1, 256 * 6 * 6)
        x = F.dropout(F.relu(self.fc1(x)))
        x = F.dropout(F.relu(self.fc2(x)))
        x = self.fc3(x)
        return x
复制代码

在本次垃圾分类任务中,最终将图片分为六类,因此与原始AlexNet不同,最后一层全连接层的输出size为6。

根据AlexNet论文中的参数,优化器使用SGD, 并将其学习率设置为0.01, 动量衰减参数设置为0.9,权重衰减参数为0.0005。 损失函数使用CrossEntropyLoss

optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=0.0005)

criterion = nn.CrossEntropyLoss()
复制代码

定义训练过程

def train(dataloader):
    epoch_loss = 0.0
    iter_num = 0

    correct = 0
    total = 0

    for i, data in enumerate(dataloader, 0):
        inputs, labels = data
        if use_gpu:
            inputs = inputs.to(GPU)
            labels = labels.to(GPU)

        if torch.is_grad_enabled():
            optimizer.zero_grad()

        outputs = net(inputs)
        
        loss = criterion(outputs, labels)

        if torch.is_grad_enabled():
            loss.backward()

            optimizer.step()

        epoch_loss += loss.item()
        iter_num += 1

        _, predicted = torch.max(outputs, 1)
        c = (predicted == labels).squeeze()
        for i, lb in enumerate(labels):
            correct += c[i].item()
            total += 1

    return epoch_loss / iter_num, correct / total
复制代码

训练模型

for epoch in range(0, EPOCH_NUMBER):
    t_l, t_a = train(trainloader)
    train_loss.append(t_l)
    train_accuracy.append(t_a)

    with torch.no_grad():
        v_l, v_a = train(validloader)
        
    print("Epoch %03d train loss: %.6f" % (epoch + 1, t_l))
    print("        val accuracy: %.2f%%" % (100 * v_a))
        
    val_loss.append(v_l)
    val_accuracy.append(v_a)
复制代码

可视化训练结果

plt.figure(figsize=(15, 5))
plt.subplot(121)
plt.plot(range(EPOCH_NUMBER), train_accuracy, label="train")
plt.plot(range(EPOCH_NUMBER), val_accuracy, label='val')
plt.title("Accuracy", size=15)
plt.legend()
plt.grid(True)
plt.subplot(122)
plt.plot(range(EPOCH_NUMBER), train_loss, label="train")
plt.plot(range(EPOCH_NUMBER), val_loss, label="val")
plt.title("Loss", size=15)
plt.legend()
plt.grid(True)
plt.show()
复制代码

从图中可以看出, 随着迭代次数的增加,准确率逐渐增加,当迭代次数超过75次之后,趋向于稳定。 在验证集上的精度可以到95%以上,与训练集差别很小,说明分类效果良好,模型泛化能力不错。

Mo项目链接