Tensorflow 中的梯度重新打包是什么?

What is gradient repacking in Tensorflow?

当 运行 来自终端的 tensorflow 基准测试时,我们可以指定几个参数。有一个参数叫做gradient_repacking。它代表什么以及如何设置它?

python tf_cnn_benchmarks.py --data_format=NCHW --batch_size=256 \
--model=resnet50 --optimizer=momentum --variable_update=replicated \
--nodistortions --gradient_repacking=8 --num_gpus=8 \
--num_epochs=90 --weight_decay=1e-4 --data_dir=${DATA_DIR} --use_fp16 \
--train_dir=${CKPT_DIR}

对于那些以后搜索的人来说,gradient_repacking 在复制模式下影响 all-reduce。从标志定义:

flags.DEFINE_integer('gradient_repacking', 0, 'Use gradient repacking. It'
                     'currently only works with replicated mode. At the end of'
                     'of each step, it repacks the gradients for more efficient'
                     'cross-device transportation. A non-zero value specifies'
                     'the number of split packs that will be formed.',
                     lower_bound=0)

至于最佳,我已经看到了 gradient_repacking=8gradient_repacking=2

我最好的猜测是参数指的是梯度被分解为在其他工作人员之间共享的碎片数量。在这种情况下,八似乎意味着每个 GPU 彼此共享(即 all-to-all)(对于您的 num_gpus=8),而 2 则意味着仅以环形方式与邻居共享。

鉴于Horovod使用自己的all reduce算法,设置gradient_repacking--variable_update=horovod时没有效果是有道理的。