content: # -*- coding: utf-8 -*-
  # File   : replicate.py
  # Author : Jiayuan Mao
  # Email  : maojiayuan@gmail.com
  # Date   : 27/01/2018
  #
  # This file is part of Synchronized-BatchNorm-PyTorch.
  # https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
  # Distributed under MIT License.
  
  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