|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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
|
|
|
