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