当我们在谈论 batch gradient descent 时，我们在谈论什么？

炼丹
deep learning, machine learning, segmentation

intuition

在 stackexchange 上看到的一个 discussion，里面提到在 deep learning 这本书中有提到一个观点，所谓的 batch gradient decent 其实只不过是一个 large mini-batch, 真实的 loss 是针对于所有样本 loss 的期望。

little math and statistics

上述的两个论断都基于一个假设，那就是我们的样本 feature 取值也是某个随机变量产生的（详见数理统计随机变量和统计量），但是通常的一种情况是这样的，我们总会把 ml/dl 中的 feature 看作常数/定值，loss 函数也无非就是在所有样本上计算 loss 求一个均值。本质上，这个均值实际上就是对上文提到的期望的一个估计罢了

不妨先从一个最简单的线性回归看起（利用 gradient decent 求解参数）

现有随机变量 $Y$, $X_1$, 且 $Y = g(X_1) + \epsilon$, 其中考虑 $g(X_1) = \beta_0 + \beta_1X_1$, 也就是线性关系假设，则有
$$
Y = \beta_0 + \beta_1X_1 + \epsilon \tag{0}
$$

我们的 objective 可犹如下的式子 formalize

$$
min\ E(Y - \beta_0 - \beta_1X_1) ^ 2 \tag{1}
$$

其中

$$
E(Y - \beta_0 - \beta_1X_1) ^ 2 \
= E_{X_1}E\left[(Y - \beta_0 - \beta_1X_1) ^ 2 | X_1\right]
$$

$E\left[(Y - \beta_0 + \beta_1X_1) ^ 2 | X_1\right]$ 为一个条件期望，在给定 $X_1 = x_1$ 的条件下，另 $y = Y|(X_1=x_1)$, 则 (1) 可重写为

$$
E(y - \beta_0 - \beta_1x_1) ^ 2 \
= \int (y - \beta_0 - \beta_1x_1) ^ 2p(x_1)dx_1 \tag{2}
$$

可以看到，给定一个 $\beta_0, \beta_1$ 真实的 loss 的形式恰如 deep learning 书中对 gradient 的论断：

The true gradient would be the expected gradient with the expectation taken over all possible examples, weighted by the data generating distribution.

而我们常用的 loss 形式只是对上述期望的一个无偏估计量罢了（均值）

$$
loss = \dfrac{\Sigma_{i=1}^n(y^{(i)} - \beta_0 - \beta_1x_1^{(i)})}{n}
$$

这也就不难理解这句话：

Using the entire training set is just using a very large minibatch size

所谓的 batch gradient descent 不过是 large mini-batch stochastic descent, 真正的 loss 是难以求得的。

batch size and lr

这一部分要说起就深了，我仅拿 facebook 的论文 Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour 来介绍一些结果，更多的讨论可见 如何评价Facebook Training ImageNet in 1 Hour这篇论文?, 以及 Tradeoff batch size vs. number of iterations to train a neural network.

在脸家的这篇论文中，他们做出 Linear Scaling Rule 的论断，在用 SGD 的前提下，如果要将 batch_size 增大的原来的 k 倍，那么 lr 也要对应增加 k 倍，并且辅以一定的 warmup iterations, 就能真正利用大 batch 的优势：训练的 iteration 少，从而更快，并且能达到小 batch 同样的精度。

one pitfall

loss 必须以最小粒度求，不能随意乘 scale, 在 gradient descent 中，loss 的 scale(比方按照整个 mini-batch 求，scale 就为 mini-batch-size) 将会影响 lr 的取值，本质上，我们可以直接认为这个 scale 是乘在 lr 之上的，这就影响了 Linear Scaling Rule。

举个例子来说，对于 semantic segmentation 问题，按照整个图片求还是按照 pixel 求影响是非常巨大的，所以一定要注意定义 loss 时，保持 最小粒度 的原则，从另一个角度看，batch_size 大训练更快的原因并不是因为 loss 变大从而 gradient 变大，而是当 batch 变大，随机性减小，网络可以更快地收敛，并且当牵扯到 batch norm 的时候，更大的 batch_size 将提供更加良好的 moving_mean/moving_variance 的估计。

one case study on deeplabv3+

1
2
3
4
5
6
7
8
9
tf.losses.softmax_cross_entropy(
    onehot_labels,
    logits,
    weights=1.0,
    label_smoothing=0,
    scope=None,
    loss_collection=tf.GraphKeys.LOSSES,
    reduction=Reduction.SUM_BY_NONZERO_WEIGHTS
)

在代码中，计算 loss 的函数为上述函数，我们仅关注 reduction=Reduction.SUM_BY_NONZERO_WEIGHTS 的讲解，其余参数请参照 api doc. 这里我去翻了翻源码 losses.softmax_cross_entropy, 该函数首先调用 nn.softmax_cross_entropy_with_logits_v2 计算 unweighted loss, 再调用 compute_weighted_loss 进行 weighted loss 的计算，reduction 仍作为一个参数传递，所以我们继续深入，在这里就可以直接看到不同 reduction 方法的内核所在，对于 Reduction.SUM_BY_NONZERO_WEIGHTS, 首先该 loss 的粒度仍是最小粒度（在这里就是每个像素），接着该方法的分母为 _num_present，也就是说，仅仅计算非零部分的最小粒度平均 loss, 而非所有 elements 的平均 loss, 这也充分 make sense, 对所有 elements 计算 loss 会算入 weights mask 掉的部分 elements, 导致 loss 偏低。

在 deeplab 官方代码中，我发现一个奇怪的现象，训练过程中我们总能看到 gpu 的个数和 loss 几乎成正比，这里是否影响了上文提到的 scaling rule 呢？

并没有，对网络参数的影响终究体现在 gradient 的计算上，所以我们可以先看 gradient, 首先不同 gpu 的 loss 在 tower_losses 的 list 中，

1
2
3
4
5
6
7
8
9
10
11
12
for i in range(FLAGS.num_clones):
  with tf.device('/gpu:%d' % i):
    # First tower has default name scope.
    name_scope = ('clone_%d' % i) if i else ''
    with tf.name_scope(name_scope) as scope:
      loss = _tower_loss(
          iterator=iterator,
          num_of_classes=num_of_classes,
          ignore_label=ignore_label,
          scope=scope,
          reuse_variable=(i != 0))
      tower_losses.append(loss)

而不同的 loss 又分别计算 gradient, 并存在 tower_grads 中

1
2
3
4
5
6
for i in range(FLAGS.num_clones):
  with tf.device('/gpu:%d' % i):
    name_scope = ('clone_%d' % i) if i else ''
    with tf.name_scope(name_scope) as scope:
      grads = optimizer.compute_gradients(tower_losses[i])
      tower_grads.append(grads)

关键点来了，_average_gradients 函数的行为是如何呢？

1
2
3
with tf.device('/cpu:0'):
    grads_and_vars = _average_gradients(tower_grads)
grad_updates = optimizer.apply_gradients(grads_and_vars, global_step=global_step)

值得一提的是，deeplab 代码使用的 synchornizaiton 策略是 parameter server 策略，有兴趣的读者可以移步 tf 分布式训练总结，而如下的 _average_gradients 恰是一个典型的 parameter server 的处理方式，在 MVSNet 的代码中也出现了同样的定义方式。通过以下的定义方式我们可以看到，尽管有多个 gpu 的 loss/gradient, 实际最终用于更新的 gradient 是 per-element 的，也就是对 gpu 求了平均，所以依然适用于上文提到的 scaling rule. 那么为什么又会出现 loss 随着 gpu 增多而增大的现象呢？

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
def _average_gradients(tower_grads):
  """Calculates average of gradient for each shared variable across all towers.
  Note that this function provides a synchronization point across all towers.
  Args:
    tower_grads: List of lists of (gradient, variable) tuples. The outer list is
      over individual gradients. The inner list is over the gradient calculation
      for each tower.
  Returns:
     List of pairs of (gradient, variable) where the gradient has been summed
       across all towers.
  """
  average_grads = []
  for grad_and_vars in zip(*tower_grads):
    # Note that each grad_and_vars looks like the following:
    #   ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
    grads, variables = zip(*grad_and_vars)
    grad = tf.reduce_mean(tf.stack(grads, axis=0), axis=0)

    # All vars are of the same value, using the first tower here.
    average_grads.append((grad, variables[0]))

  return average_grads

原因在于 tf.losses.get_total_loss() 计算 LOSS_COLLECTION 中所有 loss 的和，而这里的 total_loss 也仅仅就是 tensorboard visualization 所用，我们也可以完全手动除了 #{gpus} 来获得真正用于更新的 loss.