【论文】Alexnet实现python版本

目录结构

• data 存储训练/测试数据
  • test 测试集文件夹
  • train 训练集文件夹
  • alexnet-owt-4df8aa71.pth  alexnet 预训练数据（网上可以下载）
  • imagenet1000.json 训练数据分类标签
  • imagenet_classnames.txt 训练数据分类标签
  • tiger cat.jpg 测试用的随便找的
• results 训练结果
• src 主程序目录
  • alexnet_inference.py 用预训练数据进行对猫狗的分类测试
  • alexnet_visualizaton.py  用tensorboard包做的数据可视化，若您是pytorch1.7版本暂时有BUG无法运行
  • train_alexnet.py 训练

alexnet_inference.py

该py文件为展示模型结构、测试分类结果，下面对代码的简单介绍：

def load_class_names(p_clsnames, p_clsnames_cn):
    """
    加载标签名
    :param p_clsnames:
    :param p_clsnames_cn:
    :return:
    """
    with open(p_clsnames, "r") as f:
        class_names = json.load(f)
    with open(p_clsnames_cn, encoding='UTF-8') as f:  # 设置文件对象
        class_names_cn = f.readlines()
    return class_names, class_names_cn

load_class_names 读取分类标签的，其中标签和训练/测试图片有一一对应的关系。

def get_model(path_state_dict, vis_model=False):
    """
    创建模型，加载参数
    :param path_state_dict:
    :return:
    """
    model = models.alexnet()
    pretrained_state_dict = torch.load(path_state_dict)
    model.load_state_dict(pretrained_state_dict)
    model.eval()

    if vis_model:
        from torchsummary import summary
        summary(model, input_size=(3, 224, 224), device="cpu")

    model.to(device)
    return model

获取模型，先定义alexnet模型。导入模型预训练数据，model.eval()停止对模型数据的更新操作

def process_img(path_img):

    # hard code
    norm_mean = [0.485, 0.456, 0.406]
    norm_std = [0.229, 0.224, 0.225]
    inference_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(norm_mean, norm_std),
    ])

    # path --> img
    img_rgb = Image.open(path_img).convert('RGB')

    # img --> tensor
    img_tensor = img_transform(img_rgb, inference_transform)
    img_tensor.unsqueeze_(0)        # chw --> bchw
    img_tensor = img_tensor.to(device)

    return img_tensor, img_rgb

对读入的图片进行预处理，对图片短边裁剪,中心化(裁剪图片亦 resize(224,224))、正则化否则会效果非常差，论文中有提到。

if __name__ == "__main__":

    # config
    path_state_dict = os.path.join(BASE_DIR, "..", "data", "alexnet-owt-4df8aa71.pth")
    # path_img = os.path.join(BASE_DIR, "..", "data", "Golden Retriever from baidu.jpg")
    path_img = os.path.join(BASE_DIR, "..", "data", "tiger cat.jpg")
    path_classnames = os.path.join(BASE_DIR, "..", "data", "imagenet1000.json")
    path_classnames_cn = os.path.join(BASE_DIR, "..", "data", "imagenet_classnames.txt")

    # load class names
    cls_n, cls_n_cn = load_class_names(path_classnames, path_classnames_cn)

    # 1/5 load img
    img_tensor, img_rgb = process_img(path_img)

    # 2/5 load model
    alexnet_model = get_model(path_state_dict, True)

    # 3/5 inference  tensor --> vector
    with torch.no_grad():
        time_tic = time.time()
        outputs = alexnet_model(img_tensor)
        time_toc = time.time()

    # 4/5 index to class names
    _, pred_int = torch.max(outputs.data, 1)
    _, top5_idx = torch.topk(outputs.data, 5, dim=1)

    # 如果想把CUDA tensor格式的数据改成numpy时，需要先将其转换成cpu float-tensor随后再转到numpy格式。 numpy不能读取CUDA tensor 需要将它转化为 CPU tensor

    pred_idx = int(pred_int.cpu().numpy())
    pred_str, pred_cn = cls_n[pred_idx], cls_n_cn[pred_idx]
    print("img: {} is: {}\n{}".format(os.path.basename(path_img), pred_str, pred_cn))
    print("time consuming:{:.2f}s".format(time_toc - time_tic))

    # 5/5 visualization
    plt.imshow(img_rgb)
    plt.title("predict:{}".format(pred_str))
    top5_num = top5_idx.cpu().numpy().squeeze()
    text_str = [cls_n[t] for t in top5_num]
    for idx in range(len(top5_num)):
        plt.text(5, 15+idx*30, "top {}:{}".format(idx+1, text_str[idx]), bbox=dict(fc='yellow'))
    plt.show()

分为五个部分：

1.加载图片

2.加载模型

3.用alexnet预测

4.预测结果分别为：pred_int top1 的结果，top5_idx top5的结果

5.可视化

效果展示：

展示网络层的具体细节用到了下面的代码：

from torchsummary import summary
        summary(model, input_size=(3, 224, 224), device="cpu")

 

train_alexnet.py

 

不使用预训练参数，调大了训练模型的LR参数，直到30Epoch后模型才开始有一定的准确性，且收敛很慢。

数据库

kaggle:dogs vs cats

https://www.kaggle.com/c/dogs-vs-cats

预训练模型：

https://download.pytorch.org/models/alexnet-owt-4df8aa71.pth

源码（免费）：

https://download.csdn.net/download/qq_33000453/13210953