| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| import functools |
|
|
| from torch.nn.parallel.data_parallel import DataParallel |
|
|
| __all__ = [ |
| 'CallbackContext', |
| 'execute_replication_callbacks', |
| 'DataParallelWithCallback', |
| 'patch_replication_callback' |
| ] |
|
|
|
|
| class CallbackContext(object): |
| pass |
|
|
|
|
| def execute_replication_callbacks(modules): |
| """ |
| Execute an replication callback `__data_parallel_replicate__` on each module created by original replication. |
| |
| The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)` |
| |
| Note that, as all modules are isomorphism, we assign each sub-module with a context |
| (shared among multiple copies of this module on different devices). |
| Through this context, different copies can share some information. |
| |
| We guarantee that the callback on the master copy (the first copy) will be called ahead of calling the callback |
| of any slave copies. |
| """ |
| master_copy = modules[0] |
| nr_modules = len(list(master_copy.modules())) |
| ctxs = [CallbackContext() for _ in range(nr_modules)] |
|
|
| for i, module in enumerate(modules): |
| for j, m in enumerate(module.modules()): |
| if hasattr(m, '__data_parallel_replicate__'): |
| m.__data_parallel_replicate__(ctxs[j], i) |
|
|
|
|
| class DataParallelWithCallback(DataParallel): |
| """ |
| Data Parallel with a replication callback. |
| |
| An replication callback `__data_parallel_replicate__` of each module will be invoked after being created by |
| original `replicate` function. |
| The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)` |
| |
| Examples: |
| > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) |
| > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1]) |
| # sync_bn.__data_parallel_replicate__ will be invoked. |
| """ |
|
|
| def replicate(self, module, device_ids): |
| modules = super(DataParallelWithCallback, self).replicate(module, device_ids) |
| execute_replication_callbacks(modules) |
| return modules |
|
|
|
|
| def patch_replication_callback(data_parallel): |
| """ |
| Monkey-patch an existing `DataParallel` object. Add the replication callback. |
| Useful when you have customized `DataParallel` implementation. |
| |
| Examples: |
| > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) |
| > sync_bn = DataParallel(sync_bn, device_ids=[0, 1]) |
| > patch_replication_callback(sync_bn) |
| # this is equivalent to |
| > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) |
| > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1]) |
| """ |
|
|
| assert isinstance(data_parallel, DataParallel) |
|
|
| old_replicate = data_parallel.replicate |
|
|
| @functools.wraps(old_replicate) |
| def new_replicate(module, device_ids): |
| modules = old_replicate(module, device_ids) |
| execute_replication_callbacks(modules) |
| return modules |
|
|
| data_parallel.replicate = new_replicate |
|
|
