def _create_strategy_and_mid_level(self, optimizer_name):
strategy = self._get_strategy()
with strategy.scope():
if optimizer_name == 'sgd':
optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
elif optimizer_name == 'adagrad':
optimizer = tpu_embedding_v2_utils.Adagrad(learning_rate=0.1)
elif optimizer_name == 'adam':
optimizer = tpu_embedding_v2_utils.Adam(learning_rate=0.1)
elif optimizer_name == 'ftrl':
optimizer = tpu_embedding_v2_utils.FTRL(learning_rate=0.1)
elif optimizer_name == 'adagrad_momentum':
optimizer = tpu_embedding_v2_utils.AdagradMomentum(
learning_rate=0.1,
momentum=0.9,
use_nesterov=True,
exponent=3.0,
epsilon=0.1,
beta2=0.9)
else:
raise ValueError('optimizer is not recognized: ',
optimizer_name)
mid_level_api = self._create_mid_level(optimizer=optimizer)
return strategy, mid_level_api, optimizer
def test_optimizer_with_slot_creation_fn_non_partial(self):
def slot_creation_fn(table, slot_names, _):
slots = {}
for slot in slot_names:
# Note that we don't pass functools.partial here, so on TPU we can't
# extract the shape. We expect the error below.
slots[slot] = tf_variables.Variable(
name='{}_{}'.format(table.name, slot),
initial_value=init_ops_v2.Zeros()(shape=table.shape,
dtype=dtypes.float32),
trainable=False)
return slots
optimizer = tpu_embedding_v2_utils.Adagrad(
learning_rate=0.1,
slot_variable_creation_fn=slot_creation_fn)
strategy = self._get_strategy()
with strategy.scope():
mid_level_api = tpu_embedding_v2.TPUEmbedding(
feature_config=self.feature_config,
optimizer=optimizer)
with self.assertRaisesRegex(ValueError,
'Unable to extract initializer function'):
# We aren't going to actually run anything, so the batch_size here does
# not matter.
mid_level_api.build(self.batch_size)
def create_strategy_and_mid_level(optimizer_name):
strategy = get_strategy()
with strategy.scope():
if optimizer_name == 'sgd':
optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
elif optimizer_name == 'adagrad':
optimizer = tpu_embedding_v2_utils.Adagrad(learning_rate=0.1)
elif optimizer_name == 'adam':
optimizer = tpu_embedding_v2_utils.Adam(learning_rate=0.1)
else:
raise ValueError('optimizer is not recognized: ', optimizer_name)
embedding = create_mid_level(optimizer=optimizer)
return strategy, embedding, optimizer
def tpu_embedding_config():
feature_configs = []
for dim, vocab, name in table_data:
feature_configs.append(tpu_embedding_v2_utils.FeatureConfig(
table=tpu_embedding_v2_utils.TableConfig(
vocabulary_size=int(vocab), dim=int(dim),
initializer=init_ops_v2.Zeros(), name=name)))
optimizer = tpu_embedding_v2_utils.Adagrad(
learning_rate=0.1)
with strategy.scope():
mid_level_api = tpu_embedding_v2.TPUEmbedding(
feature_config=feature_configs,
optimizer=optimizer)
mid_level_api._output_shapes = [TensorShape(128)] * len(feature_configs)
return mid_level_api._create_config_proto()
def _create_strategy_and_mid_level(self, optimizer_name):
strategy = self._get_strategy()
with strategy.scope():
if optimizer_name == 'sgd':
optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
elif optimizer_name == 'adagrad':
optimizer = tpu_embedding_v2_utils.Adagrad(learning_rate=0.1)
elif optimizer_name == 'adam':
optimizer = tpu_embedding_v2_utils.Adam(learning_rate=0.1)
elif optimizer_name == 'ftrl':
optimizer = tpu_embedding_v2_utils.FTRL(learning_rate=0.1)
else:
raise ValueError('optimizer is not recognized: ', optimizer_name)
mid_level_api = self._create_mid_level(optimizer=optimizer)
return strategy, mid_level_api, optimizer
def test_optimizer_with_slot_creation_fn(self, use_tpu):
def slot_creation_fn(table, slot_names, _):
slots = {}
for slot in slot_names:
slots[slot] = tf_variables.Variable(
name='{}_{}'.format(table.name, slot),
initial_value=functools.partial(
init_ops_v2.Zeros(), shape=table.shape, dtype=dtypes.float32),
trainable=False)
return slots
optimizer = tpu_embedding_v2_utils.Adagrad(
learning_rate=0.1,
slot_variable_creation_fn=slot_creation_fn)
if use_tpu:
strategy = self._get_strategy()
else:
strategy = distribution_strategy_context.get_strategy()
with strategy.scope():
mid_level = tpu_embedding_v2.TPUEmbedding(
feature_config=self.feature_config,
optimizer=optimizer)
# We aren't going to actually run anything, so the batch_size here does
# not matter.
mid_level.build(self.batch_size)
video_accumulator = mid_level._variables['video']['accumulators']
user_accumulator = mid_level._variables['user']['accumulators']
if use_tpu:
# To check the table contents (ensure that it is zero rather than the
# normal initial accumulator value specified to in the optimizer config),
# we need to select the underlying table variable on TPU.
# We only have one shard on Forge.
video_accumulator = video_accumulator.variables[0]
user_accumulator = user_accumulator.variables[0]
self.assertAllClose(video_accumulator.numpy(),
np.zeros((self.table_video.vocabulary_size,
self.table_video.dim)))
self.assertAllClose(user_accumulator.numpy(),
np.zeros((self.table_user.vocabulary_size,
self.table_user.dim)))
def _create_strategy_and_mid_level(self, optimizer_name):
strategy = self._get_strategy()
# Keras optimizers has to be translated to embedding optimizer with slot
# variable creation fn properly populated.
with strategy.scope():
if optimizer_name == 'sgd':
optimizer = optimizer_v2.gradient_descent.SGD(
learning_rate=0.1)
embedding_optimizer = tpu_embedding_v2_utils.SGD(
learning_rate=0.1)
elif optimizer_name == 'adagrad':
optimizer = optimizer_v2.adagrad.Adagrad(learning_rate=0.1)
embedding_optimizer = tpu_embedding_v2_utils.Adagrad(
learning_rate=0.1,
slot_variable_creation_fn=self.
_get_slot_variable_creation_fn(optimizer))
elif optimizer_name == 'adam':
optimizer = optimizer_v2.adam.Adam(learning_rate=0.1)
embedding_optimizer = tpu_embedding_v2_utils.Adam(
learning_rate=0.1,
slot_variable_creation_fn=self.
_get_slot_variable_creation_fn(optimizer))
elif optimizer_name == 'ftrl':
optimizer = optimizer_v2.ftrl.Ftrl(learning_rate=0.1)
embedding_optimizer = tpu_embedding_v2_utils.FTRL(
learning_rate=0.1,
slot_variable_creation_fn=self.
_get_slot_variable_creation_fn(optimizer))
else:
raise ValueError('optimizer is not recognized: ',
optimizer_name)
mid_level_api = self._create_mid_level(
optimizer=embedding_optimizer)
return strategy, mid_level_api, optimizer
def test_optimizer_with_slot_creation_fn_non_partial(self):
def slot_creation_fn(table, slot_names):
slots = {}
for slot in slot_names:
# Note that we don't pass functools.partial here, so on TPU we can't
# extract the shape. We expect the error below.
slots[slot] = tf_variables.Variable(
name='{}_{}'.format(table.name, slot),
initial_value=init_ops_v2.Zeros()(shape=table.shape,
dtype=dtypes.float32),
trainable=False)
return slots
optimizer = tpu_embedding_v2_utils.Adagrad(
learning_rate=0.1,
slot_variable_creation_fn=slot_creation_fn)
strategy = self._get_strategy()
num_replicas = strategy.num_replicas_in_sync
with strategy.scope():
with self.assertRaisesRegex(ValueError,
'Unable to extract initializer function'):
tpu_embedding_v2.TPUEmbedding(
feature_config=self.feature_config,
batch_size=self.batch_size*num_replicas,
optimizer=optimizer)
