Aggregate the gradients. The aggregation is colocated with the variable. Args: all_grads (K x N x 2): A list of K lists. Each of the list is a list of N (grad, var) tuples. The variables have to be shared across the K lists. average (bool): do average or sum
(all_grads, average=True)
| 262 | |
| 263 | @under_name_scope('AggregateGradsColocate') |
| 264 | def aggregate_grads_colocate(all_grads, average=True): |
| 265 | """ |
| 266 | Aggregate the gradients. The aggregation is colocated with the variable. |
| 267 | |
| 268 | Args: |
| 269 | all_grads (K x N x 2): A list of K lists. Each of the list is a list of N (grad, var) tuples. |
| 270 | The variables have to be shared across the K lists. |
| 271 | average (bool): do average or sum |
| 272 | Returns: |
| 273 | (N x 2): A list of N (grad, var) tuples, where grad is averaged or summed over K. |
| 274 | """ |
| 275 | nr_tower = len(all_grads) |
| 276 | if nr_tower == 1: |
| 277 | return all_grads[0] |
| 278 | |
| 279 | def aggregate(grads): |
| 280 | if average: |
| 281 | return tf.multiply(tf.add_n(grads), 1.0 / nr_tower) |
| 282 | else: |
| 283 | return tf.add_n(grads) |
| 284 | |
| 285 | ret = [] |
| 286 | for idx, grad_and_vars in enumerate(zip(*all_grads)): |
| 287 | # Ngpu * 2 |
| 288 | v = grad_and_vars[0][1] |
| 289 | grads = [g for (g, _) in grad_and_vars] |
| 290 | with tf.device(v.device): # colocate summed grad with var |
| 291 | grad = aggregate(grads) |
| 292 | ret.append((grad, v)) |
| 293 | return ret |
| 294 | |
| 295 | |
| 296 | @under_name_scope('AllReduceNaive') |