Main entry for the master device in each forward pass. The messages were first collected from each devices (including the master device), and then an callback will be invoked to compute the message to be sent back to each devices (including the master device).
(self, master_msg)
| 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. |
| 116 | |
| 117 | """ |
| 118 | self._activated = True |
| 119 | |
| 120 | intermediates = [(0, master_msg)] |
| 121 | for i in range(self.nr_slaves): |
| 122 | intermediates.append(self._queue.get()) |
| 123 | |
| 124 | results = self._master_callback(intermediates) |
| 125 | assert results[0][0] == 0, "The first result should belongs to the master." |
| 126 | |
| 127 | for i, res in results: |
| 128 | if i == 0: |
| 129 | continue |
| 130 | self._registry[i].result.put(res) |
| 131 | |
| 132 | for i in range(self.nr_slaves): |
| 133 | assert self._queue.get() is True |
| 134 | |
| 135 | return results[0][1] |
| 136 | |
| 137 | @property |
| 138 | def nr_slaves(self): |