Data-parallel training with async update. It builds one tower on each GPU with shared variable scope. Every tower computes the gradients and independently applies them to the variables, without synchronizing and averaging across towers.
| 354 | |
| 355 | |
| 356 | class AsyncMultiGPUBuilder(DataParallelBuilder): |
| 357 | """ |
| 358 | Data-parallel training with async update. |
| 359 | It builds one tower on each GPU with shared variable scope. |
| 360 | Every tower computes the gradients and independently applies them to the |
| 361 | variables, without synchronizing and averaging across towers. |
| 362 | """ |
| 363 | |
| 364 | def __init__(self, towers, scale_gradient=True): |
| 365 | """ |
| 366 | Args: |
| 367 | towers(list[int]): list of GPU ids. |
| 368 | scale_gradient (bool): if True, will scale each gradient by ``1.0/nr_gpu``. |
| 369 | """ |
| 370 | super(AsyncMultiGPUBuilder, self).__init__(towers) |
| 371 | self._scale_gradient = scale_gradient |
| 372 | |
| 373 | def call_for_each_tower(self, tower_fn): |
| 374 | """ |
| 375 | Call the function `tower_fn` under :class:`TowerContext` for each tower. |
| 376 | |
| 377 | Returns: |
| 378 | a list, contains the return values of `tower_fn` on each tower. |
| 379 | """ |
| 380 | ps_device = 'cpu' if len(self.towers) >= 4 else 'gpu' |
| 381 | |
| 382 | raw_devices = ['/gpu:{}'.format(k) for k in self.towers] |
| 383 | if ps_device == 'gpu': |
| 384 | devices = [LeastLoadedDeviceSetter(d, raw_devices) for d in raw_devices] |
| 385 | else: |
| 386 | devices = [tf.train.replica_device_setter( |
| 387 | worker_device=d, ps_device='/cpu:0', ps_tasks=1) for d in raw_devices] |
| 388 | |
| 389 | return DataParallelBuilder.build_on_towers(self.towers, tower_fn, devices) |
| 390 | |
| 391 | def build(self, grad_list, get_opt_fn): |
| 392 | """ |
| 393 | Args: |
| 394 | grad_list ([[(grad, var), ...], ...]): #GPU lists to be reduced. Each is the gradients computed on each GPU. |
| 395 | get_opt_fn (-> tf.train.Optimizer): callable which returns an optimizer |
| 396 | |
| 397 | Returns: |
| 398 | tf.Operation: the training op |
| 399 | """ |
| 400 | assert len(grad_list) == len(self.towers) |
| 401 | DataParallelBuilder._check_grad_list(grad_list) |
| 402 | |
| 403 | if self._scale_gradient and len(self.towers) > 1: |
| 404 | # pretend to average the grads, in order to make async and |
| 405 | # sync have consistent effective learning rate |
| 406 | gradproc = ScaleGradient(('.*', 1.0 / len(self.towers)), verbose=False) |
| 407 | grad_list = [gradproc.process(gv) for gv in grad_list] |
| 408 | # Ngpu x Nvar x 2 |
| 409 | |
| 410 | train_ops = [] |
| 411 | opt = get_opt_fn() |
| 412 | with tf.name_scope('async_apply_gradients'): |
| 413 | for i, grad_and_vars in enumerate(zip(*grad_list)): |