An abstract `SyncMaster` object. - During the replication, as the data parallel will trigger an callback of each module, all slave devices should call `register(id)` and obtain an `SlavePipe` to communicate with the master. - During the forward pass, master device invokes `run_master`,
| 54 | |
| 55 | |
| 56 | class SyncMaster(object): |
| 57 | """An abstract `SyncMaster` object. |
| 58 | |
| 59 | - During the replication, as the data parallel will trigger an callback of each module, all slave devices should |
| 60 | call `register(id)` and obtain an `SlavePipe` to communicate with the master. |
| 61 | - During the forward pass, master device invokes `run_master`, all messages from slave devices will be collected, |
| 62 | and passed to a registered callback. |
| 63 | - After receiving the messages, the master device should gather the information and determine to message passed |
| 64 | back to each slave devices. |
| 65 | """ |
| 66 | |
| 67 | def __init__(self, master_callback): |
| 68 | """ |
| 69 | |
| 70 | Args: |
| 71 | master_callback: a callback to be invoked after having collected messages from slave devices. |
| 72 | """ |
| 73 | self._master_callback = master_callback |
| 74 | self._queue = queue.Queue() |
| 75 | self._registry = collections.OrderedDict() |
| 76 | self._activated = False |
| 77 | |
| 78 | def register_slave(self, identifier): |
| 79 | """ |
| 80 | Register an slave device. |
| 81 | |
| 82 | Args: |
| 83 | identifier: an identifier, usually is the device id. |
| 84 | |
| 85 | Returns: a `SlavePipe` object which can be used to communicate with the master device. |
| 86 | |
| 87 | """ |
| 88 | if self._activated: |
| 89 | assert self._queue.empty(), 'Queue is not clean before next initialization.' |
| 90 | self._activated = False |
| 91 | self._registry.clear() |
| 92 | future = FutureResult() |
| 93 | self._registry[identifier] = _MasterRegistry(future) |
| 94 | return SlavePipe(identifier, self._queue, future) |
| 95 | |
| 96 | def run_master(self, master_msg): |
| 97 | """ |
| 98 | Main entry for the master device in each forward pass. |
| 99 | The messages were first collected from each devices (including the master device), and then |
| 100 | an callback will be invoked to compute the message to be sent back to each devices |
| 101 | (including the master device). |
| 102 | |
| 103 | Args: |
| 104 | master_msg: the message that the master want to send to itself. This will be placed as the first |
| 105 | message when calling `master_callback`. For detailed usage, see `_SynchronizedBatchNorm` for an example. |
| 106 | |
| 107 | Returns: the message to be sent back to the master device. |
| 108 | |
| 109 | """ |
| 110 | self._activated = True |
| 111 | |
| 112 | intermediates = [(0, master_msg)] |
| 113 | for i in range(self.nr_slaves): |