def test_push_gradient_async_update(self): self.create_default_server_and_stub() self.push_gradient_test_setup() # Test applying gradients to embedding and non-embedding parameters req = elasticdl_pb2.PushGradientsRequest() for g, name in zip(self.grad_values0, self.var_names): serialize_ndarray(g, req.gradients.dense_parameters[name]) serialize_indexed_slices( self.embedding_grads0, req.gradients.embedding_tables[self._embedding_info.name], ) res = self._stub.push_gradients(req) self.assertEqual(res.accepted, True) self.assertEqual(res.version, 1) expected_values = [ v - self._lr * g for v, g in zip(self.var_values, self.grad_values0) ] for name, expected_value in zip(self.var_names, expected_values): self.assertTrue( np.allclose( expected_value, self._parameters.non_embedding_params[name].numpy(), ) ) expected_embed_table = np.copy(self.embedding_table) for gv, gi in zip( self.embedding_grads0.values, self.embedding_grads0.indices ): expected_embed_table[gi] -= self._lr * gv actual_embed_table = self._parameters.get_embedding_param( self._embedding_info.name, range(len(expected_embed_table)) ) self.assertTrue(np.allclose(expected_embed_table, actual_embed_table)) # Test applying gradients with same name for name, var in zip(self.var_names, self.var_values): self._parameters.non_embedding_params[name] = tf.Variable(var) req = elasticdl_pb2.PushGradientsRequest() serialize_ndarray( self.grad_values1[1], req.gradients.dense_parameters[self.var_names[0]], ) res = self._stub.push_gradients(req) self.assertEqual(res.accepted, True) self.assertEqual(res.version, 2) expected_values = [ self.var_values[0] - self._lr * self.grad_values1[1], self.var_values[1], ] for expected_value, name in zip(expected_values, self.var_names): self.assertTrue( np.allclose( expected_value, self._parameters.non_embedding_params[name].numpy(), ) )
def setUp(self): self.params = Parameters() self.model_pb = Model() self.infos_pb = self.model_pb.embedding_table_infos self.tensors_pb = self.model_pb.dense_parameters self.embedding_tables_pb = self.model_pb.embedding_tables self.embedding_table_name = "embedding_1" self.embedding_dim = 10 embedding_pb = self.infos_pb.add() embedding_pb.name = self.embedding_table_name embedding_pb.dim = self.embedding_dim embedding_pb.initializer = "uniform" arr1 = np.random.uniform(size=(3, 4)) serialize_ndarray(arr1, self.tensors_pb["x"]) arr2 = np.random.uniform(size=(4, 5)) serialize_ndarray(arr2, self.tensors_pb["y"]) embedding_vectors = np.random.uniform(size=(2, 10)) embedding_indices = np.array([0, 8]) serialize_indexed_slices( Tensor(None, embedding_vectors, embedding_indices), self.embedding_tables_pb[self.embedding_table_name], )
def test_pull_dense_parameters(self): self.create_default_server_and_stub() param0 = { "v0": np.random.rand(3, 2).astype(np.float32), "v1": np.random.rand(10, 32).astype(np.float32), } pull_req = elasticdl_pb2.PullDenseParametersRequest() pull_req.version = -1 # try to pull variable res = self._stub.pull_dense_parameters(pull_req) # not initialized self.assertFalse(res.initialized) # init variable req = elasticdl_pb2.Model() req.version = 1 for name, var in param0.items(): serialize_ndarray(var, req.dense_parameters[name]) res = self._stub.push_model(req) self.assertEqual(res, empty_pb2.Empty()) # pull variable back res = self._stub.pull_dense_parameters(pull_req) self.assertTrue(res.initialized) self.assertEqual(res.version, req.version) for name, pb in res.dense_parameters.items(): tensor = pb_to_ndarray(pb) self.assertTrue(np.allclose(param0[name], tensor)) # pull variable again, no param as no updated version pull_req.version = res.version res = self._stub.pull_dense_parameters(pull_req) self.assertTrue(res.initialized) self.assertEqual(res.version, pull_req.version) self.assertTrue(not res.dense_parameters)
def pull_embedding_vectors(self, request, _): result = tensor_pb2.TensorProto() if not request.ids: return result embedding_vectors = self._parameters.get_embedding_param( request.name, request.ids) serialize_ndarray(embedding_vectors, result) return result
def report_variable_to_ps(self, ps_id): model = elasticdl_pb2.Model() model.version = self._model_versions_from_ps[ps_id] if ps_id in self._ps_vars: vars = self._ps_vars[ps_id] for var in vars: serialize_ndarray(var.numpy(), model.dense_parameters[var.name]) self._ps_stubs[ps_id].push_model(model)
def report_evaluation_metrics(self, model_outputs, labels): """ report evaluation metrics to ps. """ req = elasticdl_pb2.ReportEvaluationMetricsRequest() for name, output in model_outputs.items(): output = np.concatenate(output) serialize_ndarray(output, req.model_outputs[name]) labels = np.concatenate(labels) serialize_ndarray(labels, req.labels) req.worker_id = self._worker_id self._stub.report_evaluation_metrics(req)
def push_dense_parameters(self, parameters, ps_id, version): """ Push dense parameters to PS Args: parameters: a list of Tensors ps_id: PS id version: model version """ model = elasticdl_pb2.Model() model.version = version for p in parameters: if self.parameter_to_ps[p.name] == ps_id: serialize_ndarray(p.values, model.dense_parameters[p.name]) self.ps_stubs[ps_id].push_model(model)
def report_evaluation_metrics(self, model_outputs, labels): """Report evaluation metrics to master. Args: model_outputs: dict the evaluation result on training. labels: numpy array the labels on training dataset. """ req = elasticdl_pb2.ReportEvaluationMetricsRequest() for name, output in model_outputs.items(): output = np.concatenate(output) serialize_ndarray(output, req.model_outputs[name]) labels = np.concatenate(labels) serialize_ndarray(labels, req.labels) req.worker_id = self._worker_id self._stub.report_evaluation_metrics(req)
def push_gradient_test_setup(self): self.var_names = ["test_1", "test_2"] self.var_values = [ np.array([10.0, 20.0, 30.0], np.float32), np.array([20.0, 40.0, 60.0], np.float32), ] self.grad_values0 = [ np.array([1.0, 2.0, 3.0], np.float32), np.array([2.0, 4.0, 6.0], np.float32), ] self.grad_values1 = [ np.array([0.0, 0.0, 7.0], np.float32), np.array([9.0, 9.0, 6.0], np.float32), ] dim = self._embedding_info.dim self.embedding_table = ( np.random.rand(4 * dim).reshape((4, dim)).astype(np.float32) ) self.embedding_grads0 = tf.IndexedSlices( values=np.random.rand(3 * dim) .reshape((3, dim)) .astype(np.float32), indices=(3, 1, 3), ) self.embedding_grads1 = tf.IndexedSlices( values=np.random.rand(3 * dim) .reshape((3, dim)) .astype(np.float32), indices=(2, 2, 3), ) push_model_req = elasticdl_pb2.Model() push_model_req.version = self._parameters.version for name, value in zip(self.var_names, self.var_values): serialize_ndarray(value, push_model_req.dense_parameters[name]) push_model_req.embedding_table_infos.append(self._embedding_info) self._stub.push_model(push_model_req) for name, var in zip(self.var_names, self.var_values): self._parameters.non_embedding_params[name] = tf.Variable(var) self._parameters.embedding_params[self._embedding_info.name].set( range(len(self.embedding_table)), self.embedding_table )
def to_model_pb(self): """ Convert all parameters including embedding and non-embedding parameters to `elasticdl_pb2.Model` which can be serialized. """ model_pb = elasticdl_pb2.Model() model_pb.version = self.version for name, var in self.non_embedding_params.items(): serialize_ndarray(var.numpy(), model_pb.dense_parameters[name]) for name, embedding_table in self.embedding_params.items(): # Slot embedding table is not weights in the model, so we don't # save it to checkpoint. if not embedding_table.is_slot: serialize_indexed_slices( embedding_table.to_indexed_slices(), model_pb.embedding_tables[name], ) embedding_info = embedding_table.to_embedding_table_info_pb() model_pb.embedding_table_infos.append(embedding_info) return model_pb
def pull_dense_parameters(self, request, _): """ Response with all non-embedding parameters if initialized. """ res = elasticdl_pb2.PullDenseParametersResponse() if not self._parameters.initialized: res.initialized = False return res # Only sync-SGD needs lock # TODO: use a read-write lock to support multiple concurrent reads if not self._use_async: self._lock.acquire() res.version = self._parameters.version # No need to send variables if the requester has the latest version. if self._parameters.version > request.version: for name, var in self._parameters.non_embedding_params.items(): serialize_ndarray(var.numpy(), res.dense_parameters[name]) if not self._use_async: self._lock.release() res.initialized = True return res
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 test_push_gradient_sync_update(self): self.create_server_and_stub( grads_to_wait=2, lr_staleness_modulation=False, use_async=False ) self.push_gradient_test_setup() req = elasticdl_pb2.PushGradientsRequest() req.gradients.version = 0 for g, name in zip(self.grad_values0, self.var_names): serialize_ndarray(g, req.gradients.dense_parameters[name]) serialize_indexed_slices( self.embedding_grads0, req.gradients.embedding_tables[self._embedding_info.name], ) res = self._stub.push_gradients(req) self.assertEqual(res.accepted, True) self.assertEqual(res.version, 0) req = elasticdl_pb2.PushGradientsRequest() req.gradients.version = 0 for g, name in zip(self.grad_values1, self.var_names): serialize_ndarray(g, req.gradients.dense_parameters[name]) serialize_indexed_slices( self.embedding_grads1, req.gradients.embedding_tables[self._embedding_info.name], ) res = self._stub.push_gradients(req) self.assertEqual(res.accepted, True) self.assertEqual(res.version, 1) req = elasticdl_pb2.PushGradientsRequest() req.gradients.version = 0 for g, name in zip(self.grad_values1, self.var_names): serialize_ndarray(g, req.gradients.dense_parameters[name]) res = self._stub.push_gradients(req) self.assertEqual(res.accepted, False) self.assertEqual(res.version, 1) expected_values = [ self.var_values[0] - self._lr * (self.grad_values0[0] + self.grad_values1[0]) / 2, self.var_values[1] - self._lr * (self.grad_values0[1] + self.grad_values1[1]) / 2, ] for expected_value, name in zip(expected_values, self.var_names): self.assertTrue( np.allclose( expected_value, self._parameters.non_embedding_params[name].numpy(), ) ) expected_embed_table = np.copy(self.embedding_table) for gv, gi in zip( self.embedding_grads0.values, self.embedding_grads0.indices ): expected_embed_table[gi] -= self._lr * gv for gv, gi in zip( self.embedding_grads1.values, self.embedding_grads1.indices ): expected_embed_table[gi] -= self._lr * gv actual_embed_table = self._parameters.get_embedding_param( self._embedding_info.name, range(len(expected_embed_table)) ) self.assertTrue(np.allclose(expected_embed_table, actual_embed_table))
def test_push_model(self): opt_func_name = "ftrl_optimizer" opt = load_module(_module_file).__dict__[opt_func_name]() opt_config = opt.get_config() slot_names = ["accumulator", "linear"] slot_init_value = { "accumulator": opt_config["initial_accumulator_value"], "linear": 0.0, } self.create_default_server_and_stub(optimizer=opt_func_name) param0 = { "v0": np.random.rand(3, 2).astype(np.float32), "v1": np.random.rand(10, 32).astype(np.float32), } param1 = { "v0": np.ones([3, 2], dtype=np.float32), "v1": np.ones([10, 32], dtype=np.float32), } models = [param0, param1] for idx, model in enumerate(models): req = elasticdl_pb2.Model() req.version = idx + 1 for name in model: serialize_ndarray(model[name], req.dense_parameters[name]) req.embedding_table_infos.append(self._embedding_info) res = self._stub.push_model(req) self.assertEqual(res, empty_pb2.Empty()) # self._parameters is initialized with the first push_model call # and the second push_model has no effect self.assertEqual(self._parameters.version, 1) for name in param0: self.assertTrue( np.allclose( param0[name], self._parameters.non_embedding_params[name].numpy(), ) ) self.assertEqual( self._embedding_info.name, self._parameters.embedding_params[ self._embedding_info.name ].name, ) self.assertEqual( self._embedding_info.dim, self._parameters.embedding_params[ self._embedding_info.name ].dim, ) self.assertEqual( tf.keras.initializers.get( self._embedding_info.initializer ).__class__, self._parameters.embedding_params[ self._embedding_info.name ].initializer.__class__, ) for slot_name in slot_names: name = get_slot_table_name( self._embedding_info.name, slot_name ) table = self._parameters.embedding_params[name] self.assertTrue(name, table.name) self.assertTrue(self._embedding_info.dim, table.dim) embedding = table.get([2]) self.assertTrue( (embedding - slot_init_value[slot_name] < 0.0001).all() )
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