def test_set_slot_to_optimizer(self):
embed_name = "test_emb"
indices = np.ndarray([2], dtype=np.int32)
embed_values = np.ndarray([2, 2], dtype=np.float32)
slot_values = {
"m": np.ndarray([2, 2], dtype=np.float32),
"v": np.ndarray([2, 2], dtype=np.float32),
}
params = Parameters()
params.embedding_params[embed_name] = EmbeddingTable(embed_name, 8)
for slot in ["m", "v"]:
slot_table_name = get_slot_table_name(embed_name, slot)
params.embedding_params[slot_table_name] = EmbeddingTable(
slot_table_name, 2, "0.0", True)
opt = Adam()
opt_wrapper = OptimizerWrapper(opt, None, params.get_embedding_param)
opt_wrapper._init_thread_local()
opt_wrapper._tls._unique_ids_all_layers[embed_name] = indices
opt_wrapper._create_embedding_variable(embed_name, embed_values)
opt_wrapper._get_slot_and_set_to_optimizer(embed_name)
self.assertEqual(len(opt._slots), 1)
opt_slots = list(opt._slots.values())[0]
self.assertEqual(sorted(opt_slots.keys()), ["m", "v"])
for name in ["m", "v"]:
self.assertTrue(
np.allclose(opt_slots[name].numpy(), slot_values[name]))
def test_delete_variables(self):
params = Parameters()
embed_layers = ["test_1", "test_2"]
slot_names = ["m", "v"]
dim = 8
for layer in embed_layers:
params.embedding_params[layer] = EmbeddingTable(layer, dim)
for slot in slot_names:
slot_key = get_slot_table_name(layer, slot)
params.embedding_params[slot_key] = EmbeddingTable(
slot_key, dim, "0.0", True)
opt = Adam()
opt_wrapper = OptimizerWrapper(opt, None, params.get_embedding_param,
params.set_embedding_param)
opt_wrapper._init_thread_local()
for name in embed_layers:
opt_wrapper._tls._unique_ids_all_layers[name] = np.ndarray(
[2], np.int32)
opt_wrapper._create_embedding_variable(
name, np.ndarray([2, dim], np.float32))
opt_wrapper._get_slot_and_set_to_optimizer(name)
self.assertTrue(len(opt._weights) == 4)
self.assertTrue(len(opt._slots) == 2)
for slot_dict in opt._slots.values():
self.assertTrue(len(slot_dict) == 2)
opt_wrapper._delete_variables()
self.assertTrue(len(opt._weights) == 0)
self.assertTrue(len(opt._slots) == 0)
def test_set_slot_values_to_variables(self):
layers = ["test-1", "test-2"]
slots = ["m", "v"]
id_num = 3
embedding_dims = {layer: 4 for layer in layers}
all_values = np.arange(48).reshape(12, 4).astype(np.float32)
slot_values = {}
offset = 0
for layer in layers:
for slot in slots:
start = offset
end = offset + id_num
slot_values.setdefault(layer,
{}).setdefault(slot,
all_values[start:end])
offset = end
opt = Adam()
opt_wrapper = OptimizerWrapper(opt, None, embedding_dims)
for layer in layers:
opt_wrapper._create_embedding_variable(layer, tf.zeros((1, 4)))
opt_wrapper._set_slot_values_to_variables(slot_values)
self.assertTrue(len(opt.weights) == 4)
for layer in layers:
slots_dict = None
for k, v in opt._slots.items():
if k.startswith(layer):
slots_dict = v
break
for slot in slots:
self.assertTrue(
(slots_dict[slot].numpy() == slot_values[layer][slot]
).all())
self.assertTrue(
(slots_dict[slot].numpy() ==
opt_wrapper._slot_variables[layer][slot].numpy()).all())
slots_dict[slot].assign(tf.ones((10, 4)))
self.assertTrue(
np.isclose(
opt_wrapper._slot_variables[layer][slot].numpy(),
1.0).all())
opt_wrapper._slot_variables[layer][slot].assign(-tf.ones((10,
4)))
self.assertTrue(
np.isclose(slots_dict[slot].numpy(), -1.0).all())
slot_values_new = {"test-1": {"m": np.zeros((3, 4), np.float32)}}
opt_wrapper._set_slot_values_to_variables(slot_values_new)
self.assertTrue(
np.isclose(opt_wrapper._slot_variables["test-1"]["m"].numpy(),
0.0).all())