def test_scatter_embedding_vector(self): vectors = np.array([[1, 2], [3, 4], [5, 6], [1, 7], [3, 9]]) indices = np.array([0, 1, 2, 3, 4]) num = 2 expected_results = {} for i, item_id in enumerate(indices): ps_id = int_to_id(item_id, num) if ps_id not in expected_results: item_list = [item_id] expected_results[ps_id] = [ np.expand_dims(vectors[i, :], axis=0), item_list, ] else: expected_results[ps_id][0] = np.concatenate( ( expected_results[ps_id][0], np.expand_dims(vectors[i, :], axis=0), ), axis=0, ) expected_results[ps_id][1].append(item_id) results = scatter_embedding_vector(vectors, indices, num) for ps_id in range(num): np.testing.assert_array_equal(results[ps_id][0], expected_results[ps_id][0]) self.assertListEqual(results[ps_id][1], expected_results[ps_id][1])
def report_gradient_to_ps(self, grads): self._timing.start_record_time("report_gradient") reqs = [ elasticdl_pb2.PushGradientsRequest() for i in range(self._ps_num) ] ps_grads = {} non_embed_vars_n = len(self._non_embed_vars) for g, v in zip( grads[:non_embed_vars_n], self._non_embed_vars.values() ): ps_id = self._var_to_ps[v.name] if ps_id not in ps_grads: ps_grads[ps_id] = {v.name: g} else: if v.name not in ps_grads[ps_id]: ps_grads[ps_id][v.name] = g else: if isinstance(g, tf.IndexedSlices): ps_grads[ps_id][v.name] = merge_indexed_slices( ps_grads[ps_id][v.name], g ) else: ps_grads[ps_id][v.name] += g for ps_id, pair in ps_grads.items(): for name, g in pair.items(): if isinstance(g, tf.IndexedSlices): v, i = deduplicate_indexed_slices(g.values, g.indices) ps_grads[ps_id][name] = tf.IndexedSlices(v, i) for ps_id in ps_grads: req = reqs[ps_id] for name, g in ps_grads[ps_id].items(): # Keras embedding layer has a dense parameter, # but an indexed slices type gradient if isinstance(g, tf.IndexedSlices): serialize_indexed_slices( Tensor(None, g.values.numpy(), g.indices.numpy()), req.gradients.embedding_tables[name], ) else: serialize_ndarray( g.numpy(), req.gradients.dense_parameters[name] ) edl_embedding_name_values = self._collect_edl_embedding_name_values() if edl_embedding_name_values: edl_embedding_grads = grads[non_embed_vars_n:] bet_number = 0 for name, embedding_and_ids in edl_embedding_name_values: bet_number += len(embedding_and_ids) if len(edl_embedding_grads) != bet_number: raise ValueError( "elasticdl.layers.embedding related gradient number %d " "does not match the number of its output tensor %d." % (len(edl_embedding_grads), bet_number) ) grad_accum_iter = 0 for name, embedding_and_ids in edl_embedding_name_values: g_values = None g_indices = None for _, ids in embedding_and_ids: grad = edl_embedding_grads[grad_accum_iter] grad_accum_iter += 1 # ElasticDL embedding layer with Sparse Gradients if isinstance(grad, tf.IndexedSlices): grad = grad.values if g_values is not None: g_values = tf.concat([g_values, grad], axis=0) g_indices = tf.concat([g_indices, ids], axis=0) else: g_values = grad g_indices = ids # Sum up the values of the duplicated indices in the # gradients. It can reduce the gradient payload of the # dense embedding. g_values, g_indices = deduplicate_indexed_slices( values=g_values, indices=g_indices ) results = scatter_embedding_vector( g_values.numpy(), g_indices.numpy(), self._ps_num ) for ps_id in results: req = reqs[ps_id] gv, gi = results[ps_id] serialize_indexed_slices( Tensor(None, gv, gi), req.gradients.embedding_tables[name], ) report_futures = [] for ps_id in range(self._ps_num): req = reqs[ps_id] req.gradients.version = self._model_versions_from_ps[ps_id] req.learning_rate = K.get_value(self._model.optimizer.lr) report_future = self._ps_stubs[ps_id].push_gradients.future(req) report_futures.append(report_future) accepted = False max_version = -1 for report_future in report_futures: res = report_future.result() if res.accepted: accepted = True if res.version > max_version: max_version = res.version self._timing.end_record_time("report_gradient") return accepted, max_version
def push_gradients( self, grads, edl_grads, learning_rate, model_versions, ): """ Push gradients to PS. There two kinds of gradients: - gradients of normal layers - sparse gradients of ElasticDL embedding layers """ reqs = [ elasticdl_pb2.PushGradientsRequest() for i in range(self.ps_num) ] ps_grads = {} # 1. handle grads for grad in grads: ps_id = self.parameter_to_ps[grad.name] if ps_id not in ps_grads: ps_grads[ps_id] = {grad.name: grad} else: if grad.name not in ps_grads[ps_id]: ps_grads[ps_id][grad.name] = grad else: if grad.indices is not None: ps_grads[ps_id][grad.name] = merge_indexed_slices( ps_grads[ps_id][grad.name], grad ) else: ps_grads[ps_id][grad.name].values += grad.values for ps_id, pair in ps_grads.items(): for name, grad in pair.items(): if grad.indices is not None: v, i = deduplicate_indexed_slices( grad.values, grad.indices ) ps_grads[ps_id][name] = Tensor(None, v, i) for ps_id in ps_grads: req = reqs[ps_id] for name, grad in ps_grads[ps_id].items(): # Keras embedding layer has a dense parameter, # but an indexed slices type gradient if grad.indices is not None: serialize_indexed_slices( Tensor(None, grad.values, grad.indices), req.gradients.embedding_tables[name], ) else: serialize_ndarray( grad.values, req.gradients.dense_parameters[name] ) # 2. handle sparse grads of elasticdl embedding layers groups = {} for grad in edl_grads: if grad.name not in groups: groups[grad.name] = grad else: groups[grad.name] = merge_indexed_slices( groups[grad.name], grad ) # Sum up the values of the duplicated indices in the # gradients. It can reduce the gradient payload of the # dense embedding. for name, grad in groups.items(): v, i = deduplicate_indexed_slices(grad.values, grad.indices) groups[name] = Tensor(None, v, i) results = scatter_embedding_vector( groups[name].values, groups[name].indices, self.ps_num ) for ps_id in results: req = reqs[ps_id] gv, gi = results[ps_id] serialize_indexed_slices( Tensor(None, gv, gi), req.gradients.embedding_tables[name], ) # 3. push gradients to PS report_futures = [] for ps_id in range(self.ps_num): req = reqs[ps_id] req.gradients.version = model_versions[ps_id] req.learning_rate = learning_rate report_future = self.ps_stubs[ps_id].push_gradients.future(req) report_futures.append(report_future) accepted = False max_version = -1 for report_future in report_futures: res = report_future.result() if res.accepted: accepted = True if res.version > max_version: max_version = res.version return accepted, max_version
def report_gradient_to_ps(self, grads): self._timing.start_record_time("report_gradient") reqs = [ elasticdl_pb2.PushGradientRequest() for i in range(self._ps_num) ] ps_grads = {} non_embed_vars_n = len(self._non_embed_vars) for g, v in zip(grads[:non_embed_vars_n], self._non_embed_vars.values()): ps_id = self._var_to_ps[v.name] if ps_id not in ps_grads: ps_grads[ps_id] = [(g, v.name)] else: ps_grads[ps_id].append((g, v.name)) for ps_id in ps_grads: req = reqs[ps_id] for g, name in ps_grads[ps_id]: emplace_tensor_pb_from_ndarray(req.gradients, g, name=name) edl_embedding_name_values = self._collect_edl_embedding_name_values() if edl_embedding_name_values: edl_embedding_grads = grads[non_embed_vars_n:] bet_number = 0 for name, embedding_and_ids in edl_embedding_name_values: bet_number += len(embedding_and_ids) if len(edl_embedding_grads) != bet_number: raise ValueError( "elasticdl.layers.embedding related gradient number %d " "does not match the number of its output tensor %d." % (len(edl_embedding_grads), bet_number)) grad_accum_iter = 0 for name, embedding_and_ids in edl_embedding_name_values: g_values = None g_indices = None for _, ids in embedding_and_ids: grad = edl_embedding_grads[grad_accum_iter] grad_accum_iter += 1 # ElasticDL embedding layer with Sparse Gradients if isinstance(grad, tf.IndexedSlices): grad = grad.values if g_values is not None: g_values = tf.concat([g_values, grad], axis=0) g_indices = tf.concat([g_indices, ids], axis=0) else: g_values = grad g_indices = ids # Sum up the values of the duplicated indices in the # gradients. It can reduce the gradient payload of the # dense embedding. g_values, g_indices = deduplicate_indexed_slices( values=g_values, indices=g_indices) results = scatter_embedding_vector(g_values.numpy(), g_indices.numpy(), self._ps_num) for ps_id in results: req = reqs[ps_id] gv, gi = results[ps_id] emplace_tensor_pb_from_ndarray(req.gradients, values=gv, indices=gi, name=name) report_futures = [] for ps_id in range(self._ps_num): req = reqs[ps_id] req.model_version = self._model_versions_from_ps[ps_id] report_future = self._ps_stubs[ps_id].push_gradient.future(req) report_futures.append(report_future) accepted = False max_version = -1 for report_future in report_futures: res = report_future.result() if res.accepted: accepted = True if res.model_version > max_version: max_version = res.model_version self._timing.end_record_time("report_gradient") return accepted, max_version
def report_gradient_to_ps(self, grads): reqs = [ elasticdl_pb2.PushGradientRequest() for i in range(len(self._ps_stubs)) ] ps_grads = {} non_embed_vars_n = len(self._non_embed_vars) for g, v in zip(grads[:non_embed_vars_n], self._non_embed_vars.values()): ps_id = self._var_to_ps[v.name] if ps_id not in ps_grads: ps_grads[ps_id] = [(g, v.name)] else: ps_grads[ps_id].append((g, v.name)) for ps_id in ps_grads: req = reqs[ps_id] for g, name in ps_grads[ps_id]: emplace_tensor_pb_from_ndarray(req.gradients, g, name=name) if self._embedding_layers: edl_embedding_grads = grads[non_embed_vars_n:] bet_number = 0 for layer in self._embedding_layers: bet_number += len(layer.embedding_and_ids) if len(edl_embedding_grads) != bet_number: raise ValueError( "elasticdl.layers.embedding related gradient number %d " "does not match the number of its output tensor %d." % (len(edl_embedding_grads), bet_number)) grad_accum_iter = 0 for layer in self._embedding_layers: g_values = None g_indices = None for _, ids in layer.embedding_and_ids: grad = edl_embedding_grads[grad_accum_iter] grad_accum_iter += 1 # ElasticDL embedding layer with Sparse Gradients if isinstance(grad, tf.IndexedSlices): grad = grad.values if g_values is not None: g_values = tf.concat([g_values, grad], axis=0) g_indices = tf.concat([g_indices, ids], axis=0) else: g_values = grad g_indices = ids results = scatter_embedding_vector(g_values.numpy(), g_indices.numpy(), len(self._ps_stubs)) for ps_id in results: req = reqs[ps_id] gv, gi = results[ps_id] emplace_tensor_pb_from_ndarray(req.gradients, values=gv, indices=gi, name=layer.name) report_futures = [] for ps_id in range(len(self._ps_stubs)): req = reqs[ps_id] req.model_version = self._model_version report_future = self._ps_stubs[ps_id].push_gradient.future(req) report_futures.append(report_future) for report_future in report_futures: res = report_future.result() # TODO: choose the last response temporarily return res.accepted, res.model_version