如何让Retinaface训练的更快更平滑

调优参数

学习率本身动态就会优化,所以就是初始参数训练出来一个最优解
在这里插入图片描述

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')

设置锁业内存和非阻塞方式

  1. pin_memory=True
    作用:将数据加载到锁页内存中(即物理内存),以便更快地将数据从主机内存传输到 GPU。
    影响:
    数据传输效率提高:因为数据在锁页内存中不会被交换到硬盘,这样可以减少从主机到 GPU 的数据传输时间,提高训练速度。
    适用条件:如果你的系统有足够的内存,开启 pin_memory 可以提高数据加载和训练的效率,特别是在使用大量数据和复杂模型时。
  2. 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张照片的原因吧

  • 3
    点赞
  • 2
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值