【技术博客】基于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_sizetrainset_size = dataset_size - validset_size - testset_sizetrainset, 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我的项目链接