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