CNN经典之AlexNet网络+PyTorch复现

一、前情说明：

写在前面的话
本系列博客在于汇总CSDN的精华帖，类似自用笔记，不做学习交流，方便以后的复习回顾，博文中的引用都注明出处，并点赞收藏原博主。
本博客大致分为两部分，第一部是转载于其他平台的关于LeNet的讲解，第二部分是自己对网络的复现，包括：数据集加载和预处理，网络构建，模型测试等

二、论文解读：

三、模型构建与结果展示：

3.1 下载数据集集：'http://download.tensorflow.org/example_images/flower_photos.tgz

3.2 数据集处理：
数据集下载后，解压到data_set文件夹，另外，新建sample，output，flower_data文件加夹，总的文件如下：

使用split_data将数据集分成训练集和验证集，比例在9：1之间，分好后的样本自动存入flower_data文件夹中，如下：

3.3 模型构建

建立model.py，构建神经网络：

import torch.nn as nn
import torch

class AlexNet(nn.Module):
    def __init__(self, num_classes=1000, init_weights=False):
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(64, 192, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(192, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
        self.classifier = nn.Sequential(
            nn.Dropout(),
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, num_classes),
        )
        if init_weights:
            self._initialize_weights()

    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, start_dim=1)
        x = self.classifier(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)

建立train.py，用来训练模型，记住，由于博主电脑性能较老旧，所以博主采用的是迁移学习的训练方式，硬件设施允许的情况下，还是可以试试从零开始训练的：

import torch
import torch.nn as nn
from torchvision import transforms, datasets
import torch.optim as optim
from model import AlexNet
import os
import json
import time
import torchvision.models.alexnet

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

data_transform = {
    "train": transforms.Compose([transforms.RandomResizedCrop(224),
                                 transforms.RandomHorizontalFlip(),
                                 transforms.ToTensor(),
                                 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]),
    "val": transforms.Compose([transforms.Resize((224, 224)),  # cannot 224, must (224, 224)
                               transforms.ToTensor(),
                               transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])}

data_root = os.path.abspath(os.path.join(os.getcwd(), "E:/pytorch/deep_learning/AlexNet"))  # get data root path
image_path = data_root + "/flower_data/"  # flower data set path
train_dataset = datasets.ImageFolder(root=image_path + "/train",
                                     transform=data_transform["train"])
train_num = len(train_dataset)

# {'daisy':0, 'dandelion':1, 'roses':2, 'sunflower':3, 'tulips':4}
flower_list = train_dataset.class_to_idx
cla_dict = dict((val, key) for key, val in flower_list.items())
# write dict into json file
json_str = json.dumps(cla_dict, indent=4)
with open('class_indices.json', 'w') as json_file:
    json_file.write(json_str)

batch_size = 32
train_loader = torch.utils.data.DataLoader(train_dataset,
                                           batch_size=batch_size, shuffle=True,
                                           num_workers=0)

validate_dataset = datasets.ImageFolder(root=image_path + "/val",
                                        transform=data_transform["val"])
val_num = len(validate_dataset)
validate_loader = torch.utils.data.DataLoader(validate_dataset,
                                              batch_size=4, shuffle=True,
                                              num_workers=0)

net = AlexNet()


# 加载预训练模型
model_weights_path = './alexnet-owt-4df8aa71.pth'
missing_keys, unexpected_keys = net.load_state_dict(torch.load(model_weights_path), strict=False)

net.to(device)

loss_function = nn.CrossEntropyLoss()
# pata = list(net.parameters())
optimizer = optim.Adam(net.parameters(), lr=0.0002)


if __name__ == '__main__':
    save_path = './AlexNet.pth'
    best_acc = 0.0
    for epoch in range(10):
        # train
        net.train()
        running_loss = 0.0
        t1 = time.perf_counter()
        for step, data in enumerate(train_loader, start=0):
            images, labels = data
            images = images.to(device)
            labels = labels.to(device)
            optimizer.zero_grad()
            outputs = net(images)
            loss = loss_function(outputs, labels)
            loss.backward()
            optimizer.step()

            # print statistics
            running_loss += loss.item()
            # print train process
            rate = (step + 1) / len(train_loader)
            a = "*" * int(rate * 50)
            b = "." * int((1 - rate) * 50)
            print("\rtrain loss: {:^3.0f}%[{}->{}]{:.3f}".format(int(rate * 100), a, b, loss), end="")
        print()
        print(time.perf_counter() - t1)

        # validate
        net.eval()
        acc = 0.0  # accumulate accurate number / epoch
        with torch.no_grad():
            for val_data in validate_loader:
                val_images, val_labels = val_data
                outputs = net(val_images.to(device))
                predict_y = torch.max(outputs, dim=1)[1]
                acc += (predict_y == val_labels.to(device)).sum().item()
            val_accurate = acc / val_num
            if val_accurate > best_acc:
                best_acc = val_accurate
                torch.save(net.state_dict(), save_path)
            print('[epoch %d] train_loss: %.3f  test_accuracy: %.3f' %
                  (epoch + 1, running_loss / step, val_accurate))

    print('Finished Training')

model_weights_path = ‘./alexnet-owt-4df8aa71.pth’
missing_keys, unexpected_keys =net.load_state_dict(torch.load(model_weights_path), strict=False)

大家将加载后的预训练模型放在根目录即可，也就是截图中的第二个：

训练后的权重保存为截图中第一个的形式。

训练过程如下：

建立predict.py文件，用来检测样本，样本可以从网上随便下载，几张图即可：

样本统一放在sample文件夹中。

运行predict.py，进行检测，结果如下：

其中，终端会显示每一张图片的检测率，使用opencv将检测率打在图片上，并保存在output文件夹中：

大致模型构建如上，小细节后期再补上~

欢迎加入深度学习交流群:696654483
里面有很多的研究生大佬和各行各业的大咖～