【技術博客】基於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項目鏈接