Alex模型的构建

1. 前言

在这篇文章中，我将以代码作为主要叙述对象，简要介绍近段时间学习搭建的AlexNet网络和训练工作，更多了解请关注公众号深度之眼，如有侵权请告知删文，具体数据集以及代码我会放到评论区的链接。

2.模型的搭建

2.1 用到的数据库

import os
import numpy as np
import torch.nn as nn
import torch
import sys
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
import torch.optim as optim
from matplotlib import pyplot as plt
import torchvision.models as models
from tools.my_dataset import CatDogDataset

其中最后一个数据集为自己创建的Dataset，下文中也提到。
2.2 Dataset的创建

import os
from torch.utils.data import Dataset
from PIL import Image
import random
random.seed(1)
class CatDogDataset(Dataset):
    def __init__(self,data_dir,mode="train",split_n=0.9,rng_seed=620, transform=None):
        self.mode = mode
        self.data_dir = data_dir
        self.split = split_n
        self.rng_seed = rng_seed
        self.data_info = self.get_info()
        self.transfrom = transform

    def __getitem__(self, indix):
        img,label = self.data_info[indix]
        img = Image.open(img).convert('RGB')
        if self.transfrom is not None:
            img = self.transfrom(img)

        return img, label

    def __len__(self):
        if len(self.data_info) == 0:
            raise Exception("{} is empty！".format(self.data_dir))

        return len(self.data_info)

    def get_info(self):
        img_names = os.listdir(self.data_dir)
        img_names = list(filter(lambda x: x.endswith('.jpg'),img_names))
        random.seed(self.rng_seed)
        random.shuffle(img_names)
        labels = list(0 if n.startswith('cat') else 1 for n in img_names)
        split = int(len(img_names) * self.split)
        if self.mode == 'train':
            img_set = img_names[:split]