Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def _do_push_pull_param(self, index, delta_weight):
if isinstance(index, (tuple, list)):
for i in range(len(index)):
byteps_declare_tensor("weight_" + str(index[i]))
byteps_push_pull(delta_weight[i], version=0, priority=-index[i],
name="weight_" + str(index[i]), is_average=False)
else:
byteps_declare_tensor("weight_" + str(index))
byteps_push_pull(delta_weight, version=0, priority=-index,
name="weight_" + str(index), is_average=False)
def _do_push_pull(self, index, grad):
if isinstance(index, (tuple, list)):
for i in range(len(index)):
byteps_declare_tensor("gradient_" + str(index[i]))
byteps_push_pull(grad[i], version=0, priority=-index[i],
name="gradient_" + str(index[i]), is_average=True)
else:
byteps_declare_tensor("gradient_" + str(index))
byteps_push_pull(grad, version=0, priority=-index,
name="gradient_" + str(index), is_average=True)
def _init_params(self):
tensors = []
for param in self._params_to_init:
if param._deferred_init:
tensors.append(param)
else:
param_arrays = param._check_and_get(param._data, list)
idx = self._param2idx[param.name]
if rank() != self.root_rank:
param_arrays[0].__imul__(0)
byteps_push_pull(param_arrays[0], version=0, priority=0,
name="parameter_" + str(idx), is_average=False)
param_arrays[0].wait_to_read()
self._params_to_init = tensors
def _do_push_pull(self, index, grad):
if isinstance(index, (tuple, list)):
for i in range(len(index)):
byteps_declare_tensor("gradient_" + str(index[i]))
byteps_push_pull(grad[i], version=0, priority=-index[i],
name="gradient_" + str(index[i]), is_average=True)
else:
byteps_declare_tensor("gradient_" + str(index))
byteps_push_pull(grad, version=0, priority=-index,
name="gradient_" + str(index), is_average=True)
def _do_push_pull_param(self, index, delta_weight):
if isinstance(index, (tuple, list)):
for i in range(len(index)):
byteps_declare_tensor("weight_" + str(index[i]))
byteps_push_pull(delta_weight[i], version=0, priority=-index[i],
name="weight_" + str(index[i]), is_average=False)
else:
byteps_declare_tensor("weight_" + str(index))
byteps_push_pull(delta_weight, version=0, priority=-index,
name="weight_" + str(index), is_average=False)
def _allreduce_grads(self):
for i, param in enumerate(self._params):
if param.grad_req != 'null':
byteps_push_pull(param.list_grad()[0], is_average=False,
name="gradient_" + str(i), priority=-i)