def _create_master_and_worker(self,
service_endpoint=None,
embedding_dims={}):
model_inst = custom_model()
master = MasterServicer(
2,
2,
tf.optimizers.SGD(0.1),
None,
init_var=model_inst.trainable_variables,
embedding_service_endpoint=service_endpoint,
embedding_dims=embedding_dims,
checkpoint_filename_for_init=None,
checkpoint_service=None,
evaluation_service=None,
)
worker = Worker(
1,
JobType.TRAINING_ONLY,
2,
_model_zoo_path,
model_def="test_module.custom_model",
channel=None,
)
worker.set_model(model_inst)
worker._stub = InProcessMaster(master)
return master, worker
def _create_master_and_worker(self,
service_endpoint=None,
embedding_dims={}):
model_inst = custom_model()
master = MasterServicer(
2,
2,
tf.optimizers.SGD(0.1),
None,
init_var=model_inst.trainable_variables,
embedding_service_endpoint=service_endpoint,
embedding_dims=embedding_dims,
checkpoint_filename_for_init=None,
checkpoint_service=None,
evaluation_service=None,
)
arguments = [
"--worker_id",
1,
"--job_type",
JobType.TRAINING_ONLY,
"--minibatch_size",
2,
"--model_zoo",
_model_zoo_path,
"--model_def",
"test_module.custom_model",
]
args = parse_worker_args(arguments)
worker = Worker(args)
worker.set_model(model_inst)
worker._stub = InProcessMaster(master)
return master, worker
def test_train_acceleration_with_embedding(self):
kv_store = MockKvStore()
model_inst = CustomModel()
master = MasterServicer(
2,
2,
tf.optimizers.SGD(0.1),
None,
init_var=model_inst.trainable_variables,
checkpoint_filename_for_init=None,
checkpoint_service=None,
evaluation_service=None,
)
arguments = [
"--worker_id",
1,
"--job_type",
JobType.TRAINING_ONLY,
"--minibatch_size",
32,
"--model_zoo",
_model_zoo_path,
"--model_def",
"embedding_test_module.EdlEmbeddingModel",
]
args = parse_worker_args(arguments)
worker = Worker(args)
worker._stub = InProcessMaster(master)
inputs_list = [
{
"f1": tf.constant([[0], [1], [2]], tf.int64),
"f2": tf.constant([[2], [1], [0]], tf.int64),
},
{
"f1": tf.constant([[3], [4], [3]], tf.int64),
"f2": tf.constant([[2], [1], [0]], tf.int64),
},
]
labels_list = [[0, 1, 0], [1, 1, 0]]
input_dim = 5
embedding_dim = 16
worker.set_model(model_inst)
# initialize kv store
for layer in model_inst.layers:
if isinstance(layer, Embedding):
name = layer.name
keys = [Embedding.get_key([name, i]) for i in range(input_dim)]
values = [
np.random.rand(embedding_dim).astype(np.float32)
for i in range(input_dim)
]
kv_store.update(keys, values)
with mock.patch.object(
EmbeddingService, "lookup_embedding", kv_store.lookup
), mock.patch.object(
EmbeddingService, "update_embedding", kv_store.update
):
worker._init_embedding_layer()
worker._run_model_call_before_training(inputs_list[0])
# run training process without tf.function
correct_grads = []
correct_ids_list = []
for features, labels in zip(inputs_list, labels_list):
loss, grads = worker.training_process_eagerly(features, labels)
correct_grads.append(grads)
ids = {}
for layer in worker._embedding_layers:
ids[layer.name] = layer.embedding_and_ids[0].batch_ids
correct_ids_list.append(ids)
worker._reset_embedding()
# run training process with tf.function
test_grads = []
test_ids_list = []
for features, labels in zip(inputs_list, labels_list):
self.assertFalse(worker._train_eagerly)
loss, grads = worker.training_process(features, labels)
test_grads.append(grads)
ids = {}
for layer in worker._embedding_layers:
ids[layer.name] = copy.deepcopy(
layer.embedding_and_ids[0].batch_ids
)
test_ids_list.append(ids)
worker._reset_embedding()
# compare the gradients
for test_g, correct_g in zip(test_grads, correct_grads):
for g1, g2 in zip(test_g, correct_g):
if isinstance(g1, tf.IndexedSlices):
self.assertTrue(np.isclose(g1.values, g2.values).all())
self.assertTrue(np.isclose(g1.indices, g2.indices).all())
else:
self.assertTrue(np.isclose(g1, g2).all())
for test_ids, correct_ids in zip(correct_ids_list, test_ids_list):
for layer_name in correct_ids.keys():
self.assertTrue(
tf.equal(test_ids[layer_name], correct_ids[layer_name])
.numpy()
.all()
)