U版YOLOv5学习笔记之训练初始化

源码链接：https://github.com/ultralytics/yolov5/blob/master/train.py

在train(hyp)函数中，进行训练前的准备

• 使用yaml加载data文件，获取存储训练、验证图片路径的txt文件路径
• 通过cfg文件，初始化网络结构model
• 判断并计算图片的尺度
• 获取模型优化参数biases、weight，选取优化函数

def train(hyp):
    epochs = opt.epochs  # 300  最大迭代次数
    batch_size = opt.batch_size  # 64  一个batch的大小
    weights = opt.weights  # initial training weights  初始化权重文件路径

    # Configure
    init_seeds(1)   # 设置random、np.random、torch随机种子
    with open(opt.data) as f:
        data_dict = yaml.load(f, Loader=yaml.FullLoader)  # model dict 加载yaml数据文件
    train_path = data_dict['train']  # 存储训练数据的txt文件路径
    test_path = data_dict['val']  # 存储验证数据的txt文件路径
    nc = 1 if opt.single_cls else int(data_dict['nc'])  # number of classes  获取类别数

    # Remove previous results  清空以前的结果
    # glob.glob(pathname)返回匹配pathname的可能为空的路径名列表
    for f in glob.glob('*_batch*.jpg') + glob.glob(results_file):
        os.remove(f)

    # Create model  初始化网络结构
    model = Model(opt.cfg).to(device)

    # assert（断言）用于判断一个表达式，在表达式条件为 false 的时候触发异常
    # 判断cfg文件中类别数和data文件中的类别数是否相等
    assert model.md['nc'] == nc, '%s nc=%g classes but %s nc=%g classes' % (opt.data, nc, opt.cfg, model.md['nc'])

    # Image sizes
    gs = int(max(model.stride))  # grid size (max stride)
    if any(x % gs != 0 for x in opt.img_size):  # img size 必须整除 max stride
        print('WARNING: --img-size %g,%g must be multiple of %s max stride %g' % (*opt.img_size, opt.cfg, gs))
    imgsz, imgsz_test = [make_divisible(x, gs) for x in opt.img_size]  # image sizes (train, test)

    # Optimizer
    nbs = 64  # nominal batch size
    accumulate = max(round(nbs / batch_size), 1)  # accumulate loss before optimizing
    hyp['weight_decay'] *= batch_size * accumulate / nbs  # scale weight_decay
    pg0, pg1, pg2 = [], [], []  # optimizer parameter groups  优化参数
    for k, v in model.named_parameters():
        if v.requires_grad:
            if '.bias' in k:
                pg2.append(v)  # biases
            elif '.weight' in k and '.bn' not in k:
                pg1.append(v)  # apply weight decay
            else:
                pg0.append(v)  # all else
    # 选取优化函数
    optimizer = optim.Adam(pg0, lr=hyp['lr0']) if opt.adam else \
        optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
    optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']})  # add pg1 with weight_decay
    optimizer.add_param_group({'params': pg2})  # add pg2 (biases)
    print('Optimizer groups: %g .bias, %g conv.weight, %g other' % (len(pg2), len(pg1), len(pg0)))
    del pg0, pg1, pg2
    ...
    ...
    ...