def __init__(self, base_optimizer: ThreeStepOptimizer, comm=None):
super().__init__(base_optimizer.executor)
if comm is None:
comm = CommunicationNetwork()
self.communication = comm
self.base_optimizer = base_optimizer
self.loss = base_optimizer.loss
def build_allreduce_gradients(self, network: Caffe2Network,
comm_network: CommunicationNetwork):
_load_custom_dll()
gradients = network.gradient()
ptr = comm_network.get_comm_numpy_ptr()
network.feed_tensor("mpi_comm",
ptr,
device_option=core.DeviceOption(caffe2_pb2.CPU))
# Copy GPU data to CPU
if network.is_cuda:
for (param_name, grad_name) in gradients:
with core.DeviceScope(self.network.device_option):
# we copy on the same device on where mpi_comm is
network.train_model.EnsureCPUOutput([grad_name],
grad_name + "_cpu")
# Invoke MPI
for (param_name, grad_name) in gradients:
grad_name_from = grad_name + "_cpu" if network.is_cuda else grad_name
with core.DeviceScope(core.DeviceOption(caffe2_pb2.CPU)):
# we use the copied tensor as input
network.train_model.DMpiAllReduceMean(
[grad_name_from, "mpi_comm"], grad_name_from)
# We have to copy back the communicated tensor if we are on the GPU
if network.is_cuda:
for (param_name, grad_name) in gradients:
with core.DeviceScope(self.network.device_option):
# we copy on the same device on where mpi_comm is
network.train_model.CopyFromCPUInput([grad_name + "_cpu"],
grad_name)
def __init__(self, optimizer: TFOptimizer, comm=None):
super().__init__(optimizer.executor, optimizer.loss)
if comm is None:
comm = CommunicationNetwork()
self.communication = comm
self.original_optimizer = optimizer
try:
import horovod.tensorflow as hvd
except ImportError:
raise ImportError('Cannot import Horovod')
hvd.init()
self.op = hvd.DistributedOptimizer(optimizer.op)
def __init__(self, optimizer: Optimizer, comm=CommunicationNetwork()):
super(Caffe2ConsistentDecentralized, self).__init__(optimizer.executor)
self.communication = comm
self.parameter_optimizer = optimizer
def __init__(self, optimizer: Optimizer, comm=CommunicationNetwork()):
super(Caffe2ConsistentParameterServer,
self).__init__(optimizer.network)
self.communication = comm
self.parameter_optimizer = optimizer