| 74 | |
| 75 | |
| 76 | class SyncMultiGPUTrainerParameterServer(SingleCostTrainer): |
| 77 | |
| 78 | __doc__ = SyncMultiGPUParameterServerBuilder.__doc__ + """ |
| 79 | |
| 80 | Attributes: |
| 81 | devices (list[int]): List of GPU ids. |
| 82 | |
| 83 | """ |
| 84 | |
| 85 | @map_arg(gpus=_int_to_range) |
| 86 | def __init__(self, gpus, ps_device=None): |
| 87 | """ |
| 88 | Args: |
| 89 | gpus ([int]): list of GPU ids. |
| 90 | ps_device: either 'gpu' or 'cpu', where variables are stored. |
| 91 | The default value is subject to change. |
| 92 | """ |
| 93 | self.devices = gpus |
| 94 | if ps_device is None: |
| 95 | ps_device = 'gpu' if len(gpus) <= 2 else 'cpu' |
| 96 | self._builder = SyncMultiGPUParameterServerBuilder(gpus, ps_device) |
| 97 | super(SyncMultiGPUTrainerParameterServer, self).__init__() |
| 98 | |
| 99 | def _setup_graph(self, input, get_cost_fn, get_opt_fn): |
| 100 | if len(self.devices) > 1: |
| 101 | assert isinstance(input, FeedfreeInput), input |
| 102 | tower_fn = self._make_get_grad_fn(input, get_cost_fn, get_opt_fn) |
| 103 | grad_list = self._builder.call_for_each_tower(tower_fn) |
| 104 | self.train_op = self._builder.build(grad_list, get_opt_fn) |
| 105 | return [] |
| 106 | |
| 107 | |
| 108 | def SyncMultiGPUTrainer(gpus): |
no outgoing calls
no test coverage detected