调优参数
学习率本身动态就会优化,所以就是初始参数训练出来一个最优解
from __future__ import print_function
import argparse
import math
import os
import optuna
import torch
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data as data
from data import WiderFaceDetection, detection_collate, preproc, cfg_mnet, cfg_re50
from layers.functions.prior_box import PriorBox
from layers.modules import MultiBoxLoss
from models.retinaface import RetinaFace
# 解析命令行参数
parser = argparse.ArgumentParser(description='Retinaface Training')
parser.add_argument('--training_dataset', default='./data/lst/train/label.txt', help='训练数据集目录')
parser.add_argument('--network', default='resnet50', help='Backbone 网络选择: mobile0.25 或 resnet50')
parser.add_argument('--num_workers', default=0, type=int, help='数据加载时的工作线程数')
parser.add_argument('--resume_net', default=None, help='重新训练时的已保存模型路径')
parser.add_argument('--resume_epoch', default=0, type=int, help='重新训练时的迭代轮数')
parser.add_argument('--save_folder', default='./weights/', help='保存检查点模型的目录')
# 解析参数
args = parser.parse_args()
# 如果 save_folder 目录不存在,则创建它
if not os.path.exists(args.save_folder):
os.mkdir(args.save_folder)
# 根据选择的网络初始化配置
cfg = None
if args.network == "mobile0.25":
cfg = cfg_mnet
elif args.network == "resnet50":
cfg = cfg_re50
# 设置 RGB 平均值、类别数、图像维度等
rgb_mean = (104, 117, 123) # BGR 顺序
num_classes = 2
img_dim = cfg['image_size']
num_gpu = cfg['ngpu']
batch_size = cfg['batch_size']
max_epoch = cfg['epoch']
gpu_train = cfg['gpu_train']
num_workers = args.num_workers
training_dataset = args.training_dataset
save_folder = args.save_folder
# 超参数优化目标函数
def objective(trial):
# 超参数搜索空间
initial_lr = trial.suggest_float('lr', 0.001, 0.005, log=True)
momentum = trial.suggest_float('momentum', 0.88, 0.92)
weight_decay = trial.suggest_float('weight_decay', 5e-4, 5e-4, log=True)
gamma = trial.suggest_float('gamma', 0.1, 0.1, log=True)
# 初始化 RetinaFace 模型
net = RetinaFace(cfg=cfg)
# 如果指定了 resume_net,加载预训练权重
if args.resume_net is not None:
state_dict = torch.load(args.resume_net)
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
head = k[:7]
if head == 'module.':
name = k[7:] # 移除 `module.`
else:
name = k
new_state_dict[name] = v
net.load_state_dict(new_state_dict)
# 如果有多个 GPU 可用,使用 DataParallel 进行并行训练
if num_gpu > 1 and gpu_train:
net = torch.nn.DataParallel(net).cuda()
else:
net = net.cuda()
cudnn.benchmark = True
# 定义优化器、损失函数和先验框
optimizer = optim.SGD(net.parameters(), lr=initial_lr, momentum=momentum, weight_decay=weight_decay)
criterion = MultiBoxLoss(num_classes, 0.35, True, 0, True, 7, 0.35, False)
priorbox = PriorBox(cfg, image_size=(img_dim, img_dim))
with torch.no_grad():
priors = priorbox.forward()
priors = priors.cuda()
# 训练函数
def train():
net.train()
epoch = 0 + args.resume_epoch
dataset = WiderFaceDetection(training_dataset, preproc(img_dim, rgb_mean))
epoch_size = math.ceil(len(dataset) / batch_size)
max_iter = max_epoch * epoch_size
stepvalues = (cfg['decay1'] * epoch_size, cfg['decay2'] * epoch_size)
step_index = 0
start_iter = args.resume_epoch * epoch_size if args.resume_epoch > 0 else 0
for iteration in range(start_iter, max_iter):
if iteration % epoch_size == 0:
batch_iterator = iter(data.DataLoader(dataset, batch_size, shuffle=True, num_workers=num_workers,
collate_fn=detection_collate))
epoch += 1
if iteration in stepvalues:
step_index += 1
lr = adjust_learning_rate(optimizer, gamma, epoch, step_index, iteration, epoch_size)
images, targets = next(batch_iterator)
images = images.cuda()
targets = [anno.cuda() for anno in targets]
out = net(images)
optimizer.zero_grad()
loss_l, loss_c, loss_landm = criterion(out, priors, targets)
loss = cfg['loc_weight'] * loss_l + loss_c + loss_landm
loss.backward()
optimizer.step()
return loss.item()
# 学习率调整函数
def adjust_learning_rate(optimizer, gamma, epoch, step_index, iteration, epoch_size):
warmup_epoch = 5
if epoch < warmup_epoch:
lr = initial_lr * (iteration + 1) / (epoch_size * warmup_epoch)
else:
lr = initial_lr * 0.5 * (1 + math.cos(math.pi * (epoch - warmup_epoch) / (max_epoch - warmup_epoch)))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
return lr
# 训练并返回损失
final_loss = train()
# 将超参数和对应的损失写入文件
with open('best.txt', 'a') as f :
f.write(f"Trial {trial.number} - Loss: {final_loss}\n")
f.write(f"lr: {initial_lr}, momentum: {momentum}, weight_decay: {weight_decay}, gamma: {gamma}\n\n")
return final_loss
if __name__ == '__main__':
# 使用 Optuna 进行超参数优化
study = optuna.create_study(direction='minimize')
study.optimize(objective, n_trials=10)
print('最佳超参数:')
print(study.best_params)
# 将最佳超参数写入文件
with open('best.txt', 'a') as f:
f.write('最佳超参数:\n')
f.write(str(study.best_params))
f.write('\n')
设置锁业内存和非阻塞方式
- pin_memory=True
作用:将数据加载到锁页内存中(即物理内存),以便更快地将数据从主机内存传输到 GPU。
影响:
数据传输效率提高:因为数据在锁页内存中不会被交换到硬盘,这样可以减少从主机到 GPU 的数据传输时间,提高训练速度。
适用条件:如果你的系统有足够的内存,开启 pin_memory 可以提高数据加载和训练的效率,特别是在使用大量数据和复杂模型时。 - non_blocking=True
作用:在将数据从主机传输到 GPU 时,采用非阻塞方式进行。这意味着数据传输不会阻塞主线程,从而可以实现异步数据传输。
影响:
更高效的数据传输:与 pin_memory 配合使用时,数据的传输更加高效,因为数据已经在物理内存中,非阻塞传输可以进一步减少数据传输时间。
GPU利用率提升:通过异步传输,GPU 可以在等待数据时更少地闲置,从而提升训练效率。
在训练过程中的实际效果
数据加载速度:pin_memory=True 和 non_blocking=True 配合使用可以显著提高数据加载的速度,因为数据传输到 GPU 的延迟会被最小化。这对于大型数据集和复杂模型尤为重要,因为它们对数据加载的需求更高。
训练效率:随着数据加载速度的提升,训练过程的整体效率也会提高。这意味着你可能会看到训练时间的缩短和训练过程的平稳性增加。
实际影响分析
pin_memory 的作用:在你的 train 函数中,设置 pin_memory=True 会导致 DataLoader 将数据加载到锁页内存中,从而提高数据加载速度。你在创建 DataLoader 时设置了 pin_memory=True,这会使得数据从 CPU 内存传输到 GPU 时更加高效。
non_blocking 的作用:在将数据从 CPU 转移到 GPU 时使用了 non_blocking=True,这使得数据传输不会阻塞主线程。结合 pin_memory,这种设置会进一步优化数据传输过程,使得 GPU 可以更高效地利用数据。
总结:这两个设置的结合使用可以显著提高数据加载和训练的效率,特别是在处理大规模数据集和复杂模型时。使用 pin_memory=True 和 non_blocking=True 可以确保数据传输尽可能快,从而使 GPU 计算资源得到更好的利用。
使用pin_memory=True
使用non-blocking后
优化后的代码
# 导入必要的库
from __future__ import print_function
import os
import torch
import torch.optim as optim
import torch.backends.cudnn as cudnn
import argparse
import torch.utils.data as data
from data import WiderFaceDetection, detection_collate, preproc, cfg_mnet, cfg_re50
from layers.modules import MultiBoxLoss
from layers.functions.prior_box import PriorBox
import time
import datetime
import math
from models.retinaface import RetinaFace
# 解析命令行参数
parser = argparse.ArgumentParser(description='Retinaface Training')
parser.add_argument('--training_dataset', default='./data/lst/train/label.txt', help='训练数据集目录')
parser.add_argument('--network', default='resnet50', help='Backbone 网络选择: mobile0.25 或 resnet50')
parser.add_argument('--num_workers', default=0, type=int, help='数据加载时的工作线程数')
parser.add_argument('--lr', '--learning-rate', default=0.0019325940605536783, type=float, help='初始学习率')
parser.add_argument('--momentum', default=0.9172822845787842, type=float, help='动量')
parser.add_argument('--resume_net', default=None, help='重新训练时的已保存模型路径')
parser.add_argument('--resume_epoch', default=0, type=int, help='重新训练时的迭代轮数')
parser.add_argument('--weight_decay', default=5e-4, type=float, help='SGD的权重衰减')
parser.add_argument('--gamma', default=0.1, type=float, help='SGD的学习率衰减系数')
parser.add_argument('--save_folder', default='./weights/', help='保存检查点模型的目录')
# 解析参数
args = parser.parse_args()
# 如果 save_folder 目录不存在,则创建它
if not os.path.exists(args.save_folder):
os.mkdir(args.save_folder)
# 根据选择的网络初始化配置
cfg = None
if args.network == "mobile0.25":
cfg = cfg_mnet
elif args.network == "resnet50":
cfg = cfg_re50
# 设置 RGB 平均值、类别数、图像维度等
rgb_mean = (104, 117, 123) # BGR 顺序
num_classes = 2
img_dim = cfg['image_size']
num_gpu = cfg['ngpu']
batch_size = cfg['batch_size']
max_epoch = cfg['epoch']
gpu_train = cfg['gpu_train']
num_workers = args.num_workers
momentum = args.momentum
weight_decay = args.weight_decay
initial_lr = args.lr
gamma = args.gamma
training_dataset = args.training_dataset
save_folder = args.save_folder
# 初始化 RetinaFace 模型
net = RetinaFace(cfg=cfg)
print("打印网络结构...")
print(net)
# 如果指定了 resume_net,加载预训练权重
if args.resume_net is not None:
print('加载预训练网络...')
state_dict = torch.load(args.resume_net)
# 创建一个新的 OrderedDict,不包含 `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
head = k[:7]
if head == 'module.':
name = k[7:] # 移除 `module.`
else:
name = k
new_state_dict[name] = v
net.load_state_dict(new_state_dict)
# 如果有多个 GPU 可用,使用 DataParallel 进行并行训练
if num_gpu > 1 and gpu_train:
net = torch.nn.DataParallel(net).cuda()
else:
net = net.cuda()
# 设置 CuDNN benchmark 以提高性能
cudnn.benchmark = True
# 定义优化器、损失函数和先验框
optimizer = optim.SGD(net.parameters(), lr=initial_lr, momentum=momentum, weight_decay=weight_decay)
criterion = MultiBoxLoss(num_classes, 0.35, True, 0, True, 7, 0.35, False)
priorbox = PriorBox(cfg, image_size=(img_dim, img_dim))
with torch.no_grad():
priors = priorbox.forward()
priors = priors.cuda(non_blocking=True)
# 训练函数
def train():
net.train()
epoch = 0 + args.resume_epoch
print('加载数据集...')
# 加载用于训练的 WiderFace 数据集
dataset = WiderFaceDetection(training_dataset, preproc(img_dim, rgb_mean))
# 计算每轮的迭代次数和最大迭代次数
epoch_size = math.ceil(len(dataset) / batch_size)
max_iter = max_epoch * epoch_size
stepvalues = (cfg['decay1'] * epoch_size, cfg['decay2'] * epoch_size)
step_index = 0
# 根据 resume_epoch 设置开始迭代的位置
if args.resume_epoch > 0:
start_iter = args.resume_epoch * epoch_size
else:
start_iter = 0
for iteration in range(start_iter, max_iter):
if iteration % epoch_size == 0:
# 创建批次迭代器
batch_iterator = iter(data.DataLoader(dataset, batch_size, shuffle=True, num_workers=num_workers,
collate_fn=detection_collate,pin_memory=True))
if (epoch % 10 == 0 and epoch > 0) or (epoch % 5 == 0 and epoch > cfg['decay1']):
torch.save(net.state_dict(), save_folder + cfg['name'] + '_epoch_' + str(epoch) + '.pth')
epoch += 1
load_t0 = time.time()
if iteration in stepvalues:
step_index += 1
lr = adjust_learning_rate(optimizer, gamma, epoch, step_index, iteration, epoch_size)
# 加载训练数据
images, targets = next(batch_iterator)
images = images.cuda()
targets = [anno.cuda() for anno in targets]
# 前向传播
out = net(images)
# 反向传播
optimizer.zero_grad()
loss_l, loss_c, loss_landm = criterion(out, priors, targets)
loss = cfg['loc_weight'] * loss_l + loss_c + loss_landm
loss.backward()
optimizer.step()
load_t1 = time.time()
batch_time = load_t1 - load_t0
eta = int(batch_time * (max_iter - iteration))
print('Epoch:{}/{} || Epochiter: {}/{} || Iter: {}/{} || Loc: {:.4f} Cla: {:.4f} Landm: {:.4f} || LR: {:.8f} || Batchtime: {:.4f} s || ETA: {}'
.format(epoch, max_epoch, (iteration % epoch_size) + 1,
epoch_size, iteration + 1, max_iter, loss_l.item(), loss_c.item(), loss_landm.item(), lr,
batch_time, str(datetime.timedelta(seconds=eta))))
torch.save(net.state_dict(), save_folder + cfg['name'] + '_Final.pth')
# 学习率调整函数
def adjust_learning_rate(optimizer, gamma, epoch, step_index, iteration, epoch_size):
warmup_epoch = 5
if epoch < warmup_epoch:
lr = initial_lr * (iteration + 1) / (epoch_size * warmup_epoch)
else:
lr = initial_lr * 0.5 * (1 + math.cos(math.pi * (epoch - warmup_epoch) / (max_epoch - warmup_epoch)))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
return lr
if __name__ == '__main__':
train()
结论
没看出来加这些参数有啥用,或许是因为数据量只有200张照片的原因吧