一、准备
这里使用的是Kaggle的猫狗数据集,我们先将下载的数据集分为训练集合测试集,训练集要分为cat和dog两类,里面是打好标注的图片。
二、训练
这里我使用AlexNet和ResNet-18两个模型分别训练,比较它们的效果
首先是写一个model.py文件分别定义两个模型
#定义AlexNet模型
class AlexNet(nn.Module):
def __init__(self, num_classes=1000, init_weights=False):
super(AlexNet, self).__init__()
# 用nn.Sequential()将网络打包成一个模块,精简代码
self.features = nn.Sequential( # 卷积层提取图像特征
nn.Conv2d(3, 48, kernel_size=11, stride=4, padding=2), # input[3, 224, 224] output[48, 55, 55]
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2), # output[48, 27, 27]
nn.Conv2d(48, 128, kernel_size=5, padding=2), # output[128, 27, 27]
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2), # output[128, 13, 13]
nn.Conv2d(128, 192, kernel_size=3, padding=1), # output[192, 13, 13]
nn.ReLU(inplace=True),
nn.Conv2d(192, 192, kernel_size=3, padding=1), # output[192, 13, 13]
nn.ReLU(inplace=True),
nn.Conv2d(192, 128, kernel_size=3, padding=1), # output[128, 13, 13]
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2), # output[128, 6, 6]
)
self.classifier = nn.Sequential( # 全连接层对图像分类
nn.Dropout(p=0.5), # Dropout 随机失活神经元,默认比例为0.5
nn.Linear(128 * 6 * 6, 2048),
nn.ReLU(inplace=True),
nn.Dropout(p=0.5),
nn.Linear(2048, 2048),
nn.ReLU(inplace=True),
nn.Linear(2048, 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
# 网络权重初始化,实际上 pytorch 在构建网络时会自动初始化权重
def _initialize_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d): # 若是卷积层
nn.init.kaiming_normal_(m.weight, mode='fan_out', # 用(何)kaiming_normal_法初始化权重
nonlinearity='relu')
if m.bias is not None:
nn.init.constant_(m.bias, 0) # 初始化偏重为0
elif isinstance(m, nn.Linear): # 若是全连接层
nn.init.normal_(m.weight, 0, 0.01) # 正态分布初始化
nn.init.constant_(m.bias, 0) # 初始化偏重为0
# 定义残差块
class ResidualBlock(nn.Module):
expansion = 1
def __init__(self, in_channels, out_channels, stride=1):
super(ResidualBlock, self).__init__()
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(out_channels)
self.relu = nn.ReLU(inplace=True)
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(out_channels)
self.downsample = nn.Sequential()
if stride != 1 or in_channels != out_channels:
self.downsample = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(out_channels)
)
def forward(self, x):
identity = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
out += self.downsample(identity)
out = self.relu(out)
return out
# 定义ResNet模型
class ResNet(nn.Module):
def __init__(self, block, layers, num_classes=2):
super(ResNet, self).__init__()
self.in_channels = 64
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.layer1 = self.make_layer(block, 64, layers[0], stride=1)
self.layer2 = self.make_layer(block, 128, layers[1], stride=2)
self.layer3 = self.make_layer(block, 256, layers[2], stride=2)
self.layer4 = self.make_layer(block, 512, layers[3], stride=2)
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(512 * block.expansion, num_classes)
def make_layer(self, block, out_channels, blocks, stride=1):
layers = []
layers.append(block(self.in_channels, out_channels, stride))
self.in_channels = out_channels * block.expansion
for i in range(1, blocks):
layers.append(block(self.in_channels, out_channels))
self.in_channels = out_channels * block.expansion
return nn.Sequential(*layers)
def forward(self, x):
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.layer1(out)
out = self.layer2(out)
out = self.layer3(out)
out = self.layer4(out)
out = self.avgpool(out)
out = out.view(out.size(0), -1)
out = self.fc(out)
return out
# 构建ResNet-18模型
def resnet18(num_classes=2):
return ResNet(ResidualBlock, [2, 2, 2, 2], num_classes)
然后在数据处理阶段,需要打乱图片顺序以及读取图片名对其进行标记分类,这里可能用到split_data.py文件
import os
from shutil import copy
import random
def mkfile(file):
if not os.path.exists(file):
os.makedirs(file)
# 获取所有要分类的文件夹
file_path = "E:/yolov7/fenlei/train"
train_path = 'data/train/'
validate_path = 'data/validate/'
# 列出所有花的种类
flow_cases = [clazz for clazz in os.listdir(file_path)]
# 创建出验证集和训练集文件夹,并由类名在其目录下创建五个子目录
mkfile(train_path)
mkfile(validate_path)
for clazz in flow_cases:
mkfile(train_path + clazz)
mkfile(validate_path + clazz)
# 按照比例来进行划分, 训练集和测试集 = 9 : 1
split_rate = 0.1
# 遍历所有类别的全部图像,并按照比例分成训练集合验证集
for clazz in flow_cases:
clazz_path = file_path + '/' + clazz + '/' # 某一个类别的文件夹
images = os.listdir(clazz_path) # images 列表存储来目录下的所有图片的名称
num = len(images)
sample_images = random.sample(images, k=int(num * split_rate)) # 从images列表随机sample出k个样本
for index, image in enumerate(images):
# sample_images保存的是所有取出来的图片
if image in sample_images:
image_path = clazz_path + image
new_path = validate_path + clazz
copy(image_path, new_path)
# 其他的所有图片都保留在训练集中
else:
image_path = clazz_path + image
new_path = train_path + clazz
copy(image_path, new_path)
print(f'\r[{clazz}] processing [{index + 1} / {num}]', end="") # process bar
print()
print("processing done!")
开始训练train.py
import time
import torch
from torch import nn
from model import resnet18
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
TRAIN_ROOT = r'E:/yolov7/fenlei/DeepLearning_cases-master/AlexNet/data/train'
VALIDATE_ROOT = 'E:/yolov7/fenlei/DeepLearning_cases-master/AlexNet/data/validate'
# 进行数据的处理,定义数据转换
data_transform = {
"train": transforms.Compose([transforms.RandomResizedCrop(224), # 随机裁剪,再缩放成 224×224
transforms.RandomHorizontalFlip(p=0.5), # 水平方向随机翻转,概率为 0.5, 即一半的概率翻转, 一半的概率不翻转
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]),
"validate": 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))])}
# 加载数据集
train_dataset = ImageFolder(TRAIN_ROOT, transform=data_transform['train'])
validate_dataset = ImageFolder(VALIDATE_ROOT, transform=data_transform['validate'])
# 讲数据进行小批量处理
train_dataloader = DataLoader(train_dataset,
batch_size=16,
shuffle=True,
num_workers=0)
val_dataloader = DataLoader(validate_dataset,
batch_size=16,
shuffle=True,
num_workers=0)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = resnet18(num_classes=2).to('cuda')
#model = AlexNet(num_class=2).to('cuda')选择合适的模型
# 定义一个损失函数
loss_fn = nn.CrossEntropyLoss()
# 定义一个优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)#0.01
# 学习率每隔10轮变为原来的0.5
lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)
# 定义训练函数
def train(dataloader, model, loss_fn, optimizer):
model.train()
time_start = time.perf_counter() # 对训练一个 epoch 计时
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
# 打印训练进度(使训练过程可视化)
rate = (batch + 1) / len(dataloader) # 当前进度 = 当前step / 训练一轮epoch所需总step
a = "*" * int(rate * 50)
b = "." * int((1 - rate) * 50)
print("\r{:^3.0f}%[{}->{}]".format(int(rate * 100), a, b), end="")
print('%f s' % (time.perf_counter() - time_start))
# 返回平均的loss
train_loss = loss / n
train_acc = current / n
print('train_loss' + str(train_loss))
print('train_acc' + str(train_acc))
return train_loss, train_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
# 解决中文显示问题
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
# 定义画图函数
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')
plt.xlabel('epoch')
plt.title("训练集和验证集loss值对比图")
plt.savefig('loss_comparison.png')
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')
plt.xlabel('epoch')
plt.title("训练集和验证集acc值对比图")
plt.savefig('acc_comparison.png')
plt.show()
# 开始训练
loss_train = []
acc_train = []
loss_val = []
acc_val = []
epoch = 20
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')
matplot_loss(loss_train, loss_val)
matplot_acc(acc_train, acc_val)
print('Done!')
适当调整自己的batch_size和学习率,两个模型的效果如下:
ResNet18
AlexNet
其中ResNet18的epoch=20,AlexNet的epoch=50,所以效果略有差异。
练好的模型被保存到save_model文件夹中,可以使用以下代码进行测试
import torch
from model import AlexNet
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
TRAIN_ROOT = r'E:/yolov7/fenlei/DeepLearning_cases-master/AlexNet/data/train'
VALIDATE_ROOT = 'E:/yolov7/fenlei/DeepLearning_cases-master/AlexNet/data/validate'
# 进行数据的处理,定义数据转换
data_transform = {
"train": transforms.Compose([transforms.RandomResizedCrop(224), # 随机裁剪,再缩放成 224×224
transforms.RandomHorizontalFlip(p=0.5), # 水平方向随机翻转,概率为 0.5, 即一半的概率翻转, 一半的概率不翻转
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]),
"validate": 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))])}
# 加载数据集
train_dataset = ImageFolder(TRAIN_ROOT, transform=data_transform['train'])
validate_dataset = ImageFolder(VALIDATE_ROOT, transform=data_transform['validate'])
# 讲数据进行小批量处理
train_dataloader = DataLoader(train_dataset,
batch_size=32,
shuffle=True,
num_workers=0)
val_dataloader = DataLoader(validate_dataset,
batch_size=32,
shuffle=True,
num_workers=0)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = AlexNet(num_classes=2).to(device)
# 加载模型
model.load_state_dict(torch.load("save_model/best_model.pth", map_location='cpu'))
# 获取预测结果
classes = [
"cat",
"dog",
]
# 把张量转化为照片格式
show = ToPILImage()
# 进入到验证阶段
model.eval()
target = 0
for i in range(500):
x, y = validate_dataset[i][0], validate_dataset[i][1]
#show(x).show()
x = Variable(torch.unsqueeze(x, dim=0).float(), requires_grad=True).to(device)
x = torch.tensor(x).to(device)
with torch.no_grad():
pred = model(x)
# 用argmax获取概率最大的一个物体
predicted, actual = classes[torch.argmax(pred[0])], classes[y]
print(f'predicted:"{predicted}", Actual:"{actual}"')
if predicted == actual:
target +=1
print(fr'准确率为:"{target/500}"')
以上即为全部内容!