pytorch-yolov3实现子batch功能

pytorch-yolov3实现子batch功能

1.darknet-yolov3的子batch

前言

cfg文件中：

batch=64
subdivisions=16

batch：更新权重和偏置的基本单位

batch/subdivisions：前向传播、反向传播的基本单位

具体分析请往下看…

分析

下面以batch=64，subdivisions=16为例，并结合代码来分析它们的真实意思。

首先在训练真正开始前，会根据cfg文件来搭建网络结构，其函数为parse_network_cfg(…)。在该函数里面会调用parse_net_options(…)来读取cfg中最顶上网络参数的值，这些参数值包括网络大小，学习率，当然也包括batch和subdivisions。

void parse_net_options(list *options, network *net)
{
    net->batch = option_find_int(options, "batch",1);
    ...
    int subdivs = option_find_int(options, "subdivisions",1);
    net->time_steps = option_find_int_quiet(options, "time_steps",1);
    ...
    net->batch /= subdivs;//64/16=4
    net->batch *= net->time_steps;//RNN中的概念，目标检测不考虑(time_steps=1)
    net->subdivisions = subdivs;
... ...
}

net->batch最终值为64/16=4, 而subdiviions为16.

然后开始读取images准备开始训练，每轮读取图片的数目为下面代码所示

int imgs = net->batch * net->subdivisions * ngpus;

这个代码很有疑惑性，这里的net->batch为4， 所以一个batch读取 的图片数目还是64 (=4x16x1)。

小结一下，在darknet代码中，net->batch值是恒为cfg中的batch值 / subdivision。所以cfg中的batch是指一次性读取多少张图片，而net->batch则是被subdivision分割成的小batch。

有了这个结论，我们继续看代码中 net->batch 和subdivision到底是怎么来用的。

float train_network(network net, data d)
{
    assert(d.X.rows % net.batch == 0);//d数据的每一行代表一张样本的数据，此步骤说明该函数的运行基本单位为子batch
    int batch = net.batch;//4
    int n = d.X.rows / batch;//子batch的数量=subdivisions

    int i;
    float sum = 0;
    for(i = 0; i < n; ++i){
        get_next_batch(d, batch, i*batch, net.input, net.truth);//从d中读取batch张图片到net.input中，第一个参数d包含了一次大batch的数据,也就是net.batch*net.subdivision张图片,第二个参数batch是每次循环读取到 net.input中的数据,参与训练图片的张数，第三个参数是d中偏移量,第四个参数为网络的输入数据,第五个参数为输入数据net.input对应的标签
        float err = train_network_datum(net);//训练网路,本次训练共有net.batch张图片，训练包括一次前向传播：计算每一层网络的输出并计算cost;一次反向：计算敏感度、权重更新值、偏置更新值，适时更新权重和偏置
        sum += err;//err为loss,sum是总loss
    }
    return (float)sum/(n*batch);//算平均loss,就是一次大batch的loss了，只用于显示，不进行反向传播。
}

在 train_network_datum(net)中有行代码表明，subdivisions次训练完后才update 网络权值：

float train_network_datum(network net)
{
    // 如果使用GPU则调用gpu版的train_network_datum
#ifdef GPU
    if(gpu_index >= 0) return train_network_datum_gpu(net);
#endif
    // 更新目前已经处理的图片数量：每次处理一个batch，故直接添加l.batch
    *net.seen += net.batch;
    // 标记处于训练阶段
    net.train = 1;
    forward_network(net);
    backward_network(net);
    float error = *net.cost;
    if(((*net.seen)/net.batch)%net.subdivisions == 0) update_network(net);
    return error;
}

总结

实际上网络是net->batch（子batch）张图片进行训练（前向推理和反向传播），但是升级权值是在cfg中batch数目结束后进行的。这样在比较小的显存情况下实现大batch的训练。

2.pytorch-yolov3子batch的实现

pytorch-yolov3工程：

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    ...
    parser.add_argument("--batch_size", type=int, default=64, help="size of each image batch")
    parser.add_argument("--subdivision", type=int, default=8)
    ...

    # # Get dataloader
    dataset = ListDataset(train_path, img_size=opt.img_size, augment=True, multiscale=opt.multiscale_training)
    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=int(opt.batch_size / opt.subdivision),  # 每轮加载的batch数为一个batch的图片总数/subdivision份数
        ...
    )  
	optimizer = torch.optim.SGD(model.parameters(), lr=opt.lr, momentum=0.9)#优化器
    ...
    for epoch in range(0, opt.epochs):
        ...
        for subdivision_i, (paths, imgs, targets) in enumerate(dataloader):
            ...
            imgs = imgs.to(device)  # [n,3,416,416]
            targets = targets.to(device)
            loss, outputs, A = model(imgs, targets)  # 正向传播
            loss.backward()  # 反向传播
            nn.utils.clip_grad_norm_(model.parameters(), 20) # 梯度归一化（防止梯度爆炸）
            ...
            if (subdivision_i % opt.subdivision == 0) | (subdivision_i == len(dataloader) - 1):
                batch_i += 1
                optimizer.step()  # 更新迭代（依据与batch中累加梯度进行权重和偏置的更新，详情参考如下代码）
                optimizer.zero_grad()  # 将module中的所有模型参数的梯度初始化为0 （因为一个batch的loss关于weight的导数是所有sample的loss关于weight的导数的累加和）
                ...
            ...
        ...
    ...

3.darknet与pytorch的比较（个人理解）

1.权重的更新方式：

darknet：

梯度：遍历batch中的每张照片，对于l.delta来说，每张照片是分开存的，因此其维度会达到：l.batchl.nl.out_w*l.out_h，每个batch都会清零，batch间不会有梯度的叠加过程。

△W和△B：对于l.weight_updates以及上面提到的l.bias_updates，是将所有照片对应元素叠加起来。

update：直接用△W和△B进行更新，与梯度无关。

以卷积层为例：

void update_convolutional_layer(convolutional_layer l, int batch, float learning_rate, float momentum, float decay)
{
    int size = l.size*l.size*l.c*l.n;
    axpy_cpu(l.n, learning_rate/batch, l.bias_updates, 1, l.biases, 1);
    scal_cpu(l.n, momentum, l.bias_updates, 1);

    if(l.scales){
        axpy_cpu(l.n, learning_rate/batch, l.scale_updates, 1, l.scales, 1);
        scal_cpu(l.n, momentum, l.scale_updates, 1);
    }

    axpy_cpu(size, -decay*batch, l.weights, 1, l.weight_updates, 1);
    axpy_cpu(size, learning_rate/batch, l.weight_updates, 1, l.weights, 1);
    scal_cpu(size, momentum, l.weight_updates, 1);
}

pytorch：

梯度：一个batch的loss关于weight的导数是所有sample的loss关于weight的导数的累加和，batch间必须手动清零，否则当前batch的梯度会包含上个batch的信息。

update：没有△W和△B，而是在step()函数中直接一句梯度计算更新。

如下所示是Pytorch 中SGD优化算法的step()函数：

def step(self, closure=None):
        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:
            weight_decay = group['weight_decay']
            momentum = group['momentum']
            dampening = group['dampening']
            nesterov = group['nesterov']

            for p in group['params']:
                if p.grad is None:
                    continue
                d_p = p.grad.data
                if weight_decay != 0:
                    d_p.add_(weight_decay, p.data)
                if momentum != 0:
                    param_state = self.state[p]
                    if 'momentum_buffer' not in param_state:
                        buf = param_state['momentum_buffer'] = d_p.clone()
                    else:
                        buf = param_state['momentum_buffer']
                        buf.mul_(momentum).add_(1 - dampening, d_p)
                    if nesterov:
                        d_p = d_p.add(momentum, buf)
                    else:
                        d_p = buf

                p.data.add_(-group['lr'], d_p)#依据保存的梯度更新权重文件
        return loss

2.梯度爆炸问题：

darknet：

未找到梯度的正则化和裁剪，而且并未出现梯度的爆炸。

pytorch：

解决梯度爆炸的方法：

nn.utils.clip_grad_norm_(model.parameters(), 20) # 梯度归一化（防止梯度爆炸）

def clip_grad_norm_(parameters, max_norm, norm_type=2):
    r"""Clips gradient norm of an iterable of parameters.

    The norm is computed over all gradients together, as if they were
    concatenated into a single vector. Gradients are modified in-place.

    Arguments:
        parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
            single Tensor that will have gradients normalized
        max_norm (float or int): max norm of the gradients
        norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
            infinity norm.

    Returns:
        Total norm of the parameters (viewed as a single vector).
    """
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]
    parameters = list(filter(lambda p: p.grad is not None, parameters))
    max_norm = float(max_norm)
    norm_type = float(norm_type)
    if norm_type == inf:
        total_norm = max(p.grad.data.abs().max() for p in parameters)
    else:
        total_norm = 0
        for p in parameters:
            param_norm = p.grad.data.norm(norm_type)
            total_norm += param_norm.item() ** norm_type
        total_norm = total_norm ** (1. / norm_type)
    clip_coef = max_norm / (total_norm + 1e-6)
    if clip_coef < 1:
        for p in parameters:
            p.grad.data.mul_(clip_coef)
    return total_norm