alex搭建分类网络 pytorch

https://blog.csdn.net/didiaopao/article/details/120717525

数据准备

很多例子做图像分类的时候都喜欢用手写数字作为例子来讲解图像分类，这是一个及其不负责任的教学，我个人认为做深度学习有时候是要在数据集上下功夫的，而且因为很多框架都内置了手写数字数据集，并且已经给我们处理好了，直接可以导入到神经网络中用了，因此整个实验下来，我们连数据是什么样子都不知道，更别提学完之后去训练自己的数据集了。

这里我用的是猫狗分类的数据集，如下图所示：

利用如下脚本将数据集划分为训练集和验证集

import os
from shutil import copy
import random
  
  
def mkfile(file):
    if not os.path.exists(file):
        os.makedirs(file)
  
  
# 获取data文件夹下所有文件夹名（即需要分类的类名）
file_path = 'D:/PycharmProjects/pytorch_test/test/data_name'
flower_class = [cla for cla in os.listdir(file_path)]
  
# 创建 训练集train 文件夹，并由类名在其目录下创建5个子目录
mkfile('data/train')
for cla in flower_class:
    mkfile('data/train/' + cla)
  
# 创建 验证集val 文件夹，并由类名在其目录下创建子目录
mkfile('data/val')
for cla in flower_class:
    mkfile('data/val/' + cla)
  
# 划分比例，训练集 : 验证集 = 9 : 1
split_rate = 0.1
  
# 遍历所有类别的全部图像并按比例分成训练集和验证集
for cla in flower_class:
    cla_path = file_path + '/' + cla + '/'  # 某一类别的子目录
    images = os.listdir(cla_path)  # iamges 列表存储了该目录下所有图像的名称
    num = len(images)
    eval_index = random.sample(images, k=int(num * split_rate))  # 从images列表中随机抽取 k 个图像名称
    for index, image in enumerate(images):
        # eval_index 中保存验证集val的图像名称
        if image in eval_index:
            image_path = cla_path + image
            new_path = 'data/val/' + cla
            copy(image_path, new_path)  # 将选中的图像复制到新路径
  
        # 其余的图像保存在训练集train中
        else:
            image_path = cla_path + image
            new_path = 'data/train/' + cla
            copy(image_path, new_path)
        print("\r[{}] processing [{}/{}]".format(cla, index + 1, num), end="")  # processing bar
    print()
  
print("processing done!")

利用pytorch搭建AlexNet网络模型

import torch
from torch import nn
import torch.nn.functional as F
  
  
class MyAlexNet(nn.Module):
    def __init__(self):
        super(MyAlexNet, self).__init__()
        self.c1 = nn.Conv2d(in_channels=3, out_channels=48, kernel_size=11, stride=4, padding=2)
        self.ReLU = nn.ReLU()
        self.c2 = nn.Conv2d(in_channels=48, out_channels=128, kernel_size=5, stride=1, padding=2)
        self.s2 = nn.MaxPool2d(2)
        self.c3 = nn.Conv2d(in_channels=128, out_channels=192, kernel_size=3, stride=1, padding=1)
        self.s3 = nn.MaxPool2d(2)
        self.c4 = nn.Conv2d(in_channels=192, out_channels=192, kernel_size=3, stride=1, padding=1)
        self.c5 = nn.Conv2d(in_channels=192, out_channels=128, kernel_size=3, stride=1, padding=1)
        self.s5 = nn.MaxPool2d(kernel_size=3, stride=2)
        self.flatten = nn.Flatten()
        self.f6 = nn.Linear(4608, 2048)
        self.f7 = nn.Linear(2048, 2048)
        self.f8 = nn.Linear(2048, 1000)
  
        self.f9 = nn.Linear(1000, 2)
  
    def forward(self, x):
        x = self.ReLU(self.c1(x))
        x = self.ReLU(self.c2(x))
        x = self.s2(x)
        x = self.ReLU(self.c3(x))
        x = self.s3(x)
        x = self.ReLU(self.c4(x))
        x = self.ReLU(self.c5(x))
        x = self.s5(x)
        x = self.flatten(x)
        x = self.f6(x)
        x = F.dropout(x, p=0.5)
        x = self.f7(x)
        x = F.dropout(x, p=0.5)
        x = self.f8(x)
        x = F.dropout(x, p=0.5)
  
        x = self.f9(x)
        x = F.dropout(x, p=0.5)
        return x
  
if __name__ =="__mian__":
    x = torch.rand([1, 3, 224, 224])
    model = MyAlexNet()
    y = model(x)

训练网络模型代码

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

import torch
from torch import nn
from net import MyAlexNet
import numpy as np
  
from torch.optim import lr_scheduler
import os
  
from torchvision import transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
  
import matplotlib.pyplot as plt
  
# 解决中文显示问题
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
  
ROOT_TRAIN = r'D:/PycharmProjects/pytorch_test/test/data/train'
ROOT_TEST = r'D:/PycharmProjects/pytorch_test/test/data/val'
  
  
# 将图像RGB三个通道的像素值分别减去0.5,再除以0.5.从而将所有的像素值固定在[-1,1]范围内
normalize = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
train_transform = transforms.Compose([
    transforms.Resize((224, 224)), # 裁剪为224*224
    transforms.RandomVerticalFlip(), # 随机垂直旋转
    transforms.ToTensor(), # 将0-255范围内的像素转为0-1范围内的tensor
    normalize])
  
val_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    normalize])
  
train_dataset = ImageFolder(ROOT_TRAIN, transform=train_transform)
val_dataset = ImageFolder(ROOT_TEST, transform=val_transform)
  
train_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_dataloader = DataLoader(val_dataset, batch_size=32, shuffle=True)
  
# 如果显卡可用，则用显卡进行训练
device = 'cuda' if  torch.cuda.is_available() else 'cpu'
  
# 调用net里面的定义的网络模型， 如果GPU可用则将模型转到GPU
model = MyAlexNet().to(device)
  
# 定义损失函数（交叉熵损失）
loss_fn = nn.CrossEntropyLoss()
  
# 定义优化器（SGD）
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
  
# 学习率每隔10epoch变为原来的0.1
lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)
  
# 定义训练函数
def train(dataloader, model, loss_fn, optimizer):
    loss, current, n = 0.0, 0.0, 0
    for batch, (x, y) in enumerate(dataloader):
  
        # 前向传播
        image, y = x.to(device), y.to(device)
        output = model(image)
        cur_loss = loss_fn(output, y)
        _, pred = torch.max(output, axis=1)
        cur_acc = torch.sum(y==pred)/output.shape[0]
  
        # 反向传播
        optimizer.zero_grad()
        cur_loss.backward()
        optimizer.step()
        loss += cur_loss.item()
        current += cur_acc.item()
        n = n+1
  
    train_loss = loss / n
    tran_acc = current /n
    print('train_loss:' + str(train_loss))
    print('train_acc:' + str(tran_acc))
    return train_loss, tran_acc
  
# 定义测试函数
def val(dataloader, model, loss_fn):
    # 将模型转为验证模型
    model.eval()
    loss, current, n = 0.0, 0.0, 0
    with torch.no_grad():
        for batch, (x, y) in enumerate(dataloader):
            image, y = x.to(device), y.to(device)
            output = model(image)
            cur_loss = loss_fn(output, y)
            _, pred = torch.max(output, axis=1)
            cur_acc = torch.sum(y == pred) / output.shape[0]
            loss += cur_loss.item()
            current += cur_acc.item()
            n = n+1
  
        val_loss = loss / n
        val_acc = current / n
        print('val_loss:' + str(val_loss))
        print('val_acc:' + str(val_acc))
        return val_loss, val_acc
  
# 画图函数
def matplot_loss(train_loss, val_loss):
    plt.plot(train_loss, label='train_loss')
    plt.plot(val_loss, label='val_loss')
    plt.legend(loc='best')
    plt.ylabel('loss', fontsize=12)
    plt.xlabel('epoch', fontsize=12)
    plt.title("训练集和验证集loss值对比图")
    plt.show()
  
def matplot_acc(train_acc, val_acc):
    plt.plot(train_acc, label='train_acc')
    plt.plot(val_acc, label='val_acc')
    plt.legend(loc='best')
    plt.ylabel('acc', fontsize=12)
    plt.xlabel('epoch', fontsize=12)
    plt.title("训练集和验证集精确度值对比图")
    plt.show()
  
  
  
# 开始训练
loss_train = []
acc_train = []
loss_val = []
acc_val = []
  
epoch = 100
min_acc = 0
for t in range(epoch):
    lr_scheduler.step()
    print(f"epoch{t+1}\n--------------")
    train_loss, train_acc = train(train_dataloader, model, loss_fn, optimizer)
    val_loss, val_acc = val(val_dataloader, model, loss_fn)
  
    loss_train.append(train_loss)
    acc_train.append(train_acc)
    loss_val.append(val_loss)
    acc_val.append(val_acc)
  
    # 保存最好的模型权重文件
    if val_acc > min_acc:
        folder = 'save_model'
        if not os.path.exists(folder):
            os.mkdir('save_model')
        min_acc = val_acc
        print(f'save best model,第{t+1}轮')
        torch.save(model.state_dict(), 'save_model/best_model.pth')
    # 保存最后的权重模型文件
    if t == epoch - 1:
        torch.save(model.state_dict(), 'save_model/last_model.pth')
print('Done！')
  
matplot_loss(loss_train, loss_val)
matplot_acc(acc_train, acc_val)

测试代码

import torch
from net import MyAlexNet
import numpy as np
from torch.autograd import Variable
from torchvision import datasets, transforms
from torchvision.transforms import ToPILImage
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
  
ROOT_TRAIN = r'D:/PycharmProjects/pytorch_test/test/data/train'
ROOT_TEST = r'D:/PycharmProjects/pytorch_test/test/data/val'
  
# 将图像RGB三个通道的像素值分别减去0.5,再除以0.5.从而将所有的像素值固定在[-1,1]范围内
normalize = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
  
train_transform = transforms.Compose([
    transforms.Resize((224, 224)), # 裁剪为224*224
    transforms.RandomVerticalFlip(), # 随机垂直旋转
    transforms.ToTensor(), # 将0-255范围内的像素转为0-1范围内的tensor
    normalize])
  
val_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    normalize])
  
train_dataset = ImageFolder(ROOT_TRAIN, transform=train_transform)
val_dataset = ImageFolder(ROOT_TEST, transform=val_transform)
  
train_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_dataloader = DataLoader(val_dataset, batch_size=32, shuffle=True)
  
  
# 如果显卡可用，则用显卡进行训练
device = 'cuda' if  torch.cuda.is_available() else 'cpu'
  
# 调用net里面的定义的网络模型， 如果GPU可用则将模型转到GPU
model = MyAlexNet().to(device)
  
# 加载模型train.py里面训练的模型
model.load_state_dict(torch.load('D:/PycharmProjects/pytorch_test/test/save_model/best_model.pth'))
  
# 获取预测结果
classes = [
    'cat',
    'dag',
]
  
# 把tensor转成Image,方便可视化
show = ToPILImage()
  
# 进入验证阶段
model.eval()
  
# 对val_dataset里面的照片进行推理验证
for i in range(50):
    x, y = val_dataset[i][0], val_dataset[i][1]
    show(x).show()
    x = Variable(torch.unsqueeze(x, dim=0).float(), requires_grad=False).to(device)
    x = torch.tensor(x).to(device)
    with torch.no_grad():
        pred = model(x)
        predicted, actual = classes[torch.argmax(pred[0])], classes[y]
        print(f'Predicted: "{predicted}", Actual: "{actual}"')