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])
(data_parallel)
| 68 | |
| 69 | |
| 70 | def patch_replication_callback(data_parallel): |
| 71 | """ |
| 72 | Monkey-patch an existing `DataParallel` object. Add the replication callback. |
| 73 | Useful when you have customized `DataParallel` implementation. |
| 74 | |
| 75 | Examples: |
| 76 | > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) |
| 77 | > sync_bn = DataParallel(sync_bn, device_ids=[0, 1]) |
| 78 | > patch_replication_callback(sync_bn) |
| 79 | # this is equivalent to |
| 80 | > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) |
| 81 | > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1]) |
| 82 | """ |
| 83 | |
| 84 | assert isinstance(data_parallel, DataParallel) |
| 85 | |
| 86 | old_replicate = data_parallel.replicate |
| 87 | |
| 88 | @functools.wraps(old_replicate) |
| 89 | def new_replicate(module, device_ids): |
| 90 | modules = old_replicate(module, device_ids) |
| 91 | execute_replication_callbacks(modules) |
| 92 | return modules |
| 93 | |
| 94 | data_parallel.replicate = new_replicate |
nothing calls this directly
no outgoing calls
no test coverage detected