def _GenerateTestInputs(self):
np.random.seed(0)
weights = np.random.randn(self._num_classes, self._dim).astype(np.float32)
biases = np.random.randn(self._num_classes).astype(np.float32)
hidden_acts = np.random.randn(self._batch_size,
self._dim).astype(np.float32)
with ops.Graph().as_default() as g:
sharded_weights = variable_scope.get_variable(
"w",
partitioner=partitioned_variables.fixed_size_partitioner(
self._num_shards),
initializer=constant_op.constant(weights))
sharded_biases = variable_scope.get_variable(
"b",
partitioner=partitioned_variables.fixed_size_partitioner(
self._num_shards),
initializer=constant_op.constant(biases))
with self.test_session(graph=g) as sess:
variables.global_variables_initializer().run()
sharded_weights_v, sharded_biases_v = sess.run(
[list(sharded_weights), list(sharded_biases)])
return weights, biases, hidden_acts, sharded_weights_v, sharded_biases_v
def testMetaGraphSaveLoad(self):
save_prefix = os.path.join(self.get_temp_dir(), "ckpt")
save_graph = ops.Graph()
with save_graph.as_default(), self.test_session(
graph=save_graph) as session:
partitioner = partitioned_variables.fixed_size_partitioner(5, axis=0)
with variable_scope.variable_scope("root", partitioner=partitioner):
v0 = variable_scope.get_variable(
"v0", dtype=dtypes.float32, shape=(10, 10))
v0_list = v0._get_variable_list()
v0_part = v0._get_partitions()
self.assertEqual(len(v0_list), 5)
self.assertAllEqual(v0_part, (5, 1))
variables.global_variables_initializer().run()
save_graph.get_collection_ref("partvar").append(v0)
saver = saver_lib.Saver()
save_graph.finalize()
save_path = saver.save(sess=session, save_path=save_prefix)
previous_value = session.run(
save_graph.get_tensor_by_name(v0.name + ":0"))
restore_graph = ops.Graph()
with restore_graph.as_default(), self.test_session(
graph=restore_graph) as session:
saver = saver_lib.import_meta_graph(save_path + ".meta")
saver.restore(sess=session, save_path=save_path)
v0, = save_graph.get_collection_ref("partvar")
self.assertIsInstance(v0, variables.PartitionedVariable)
self.assertAllEqual(
previous_value,
session.run(restore_graph.get_tensor_by_name(v0.name + ":0")))
def test_load_linear_multiclass_bias_initializer(self):
"""Tests for the bias initializer wrapper."""
bias_loading_initializer = (
contrib_framework.load_linear_multiclass_bias_initializer(
new_class_vocab_file=self.new_class_vocab_file,
old_class_vocab_file=self.old_class_vocab_file,
new_class_vocab_size=4,
bias_tensor_name='some_scope/bias',
ckpt_path=[self.bundle_file],
num_class_oov_buckets=1,
initializer=self.initializer))
expected_remapped_bias_vector = np.reshape(
[2, 0, self.init_val, 1, self.init_val], [5, 1])
# The new bias vector is of size [4 class vocab + 1 class OOV, 1].
remapped_bias_vector = variable_scope.get_variable(
name='bias/obtained_bias_vector',
shape=[5, 1],
initializer=bias_loading_initializer,
partitioner=partitioned_variables.fixed_size_partitioner(3))
with self.test_session():
variables.global_variables_initializer().run()
self.assertAllClose(expected_remapped_bias_vector,
remapped_bias_vector.as_tensor().eval())
def test_load_and_remap_output_layer_weight_initializer_dnn_output(self):
"""Tests for the output layer initializer in the DNN output case."""
loading_initializer = (checkpoint_ops._load_and_remap_matrix_initializer(
new_row_vocab_size=5,
new_col_vocab_file=self.new_class_vocab_file,
old_col_vocab_file=self.old_class_vocab_file,
new_col_vocab_size=4,
old_tensor_name='some_scope/embeddings',
ckpt_path=[self.checkpoint_file],
num_col_oov_buckets=1,
initializer=self.initializer))
expected_remapped_matrix = np.concatenate(
[
np.reshape([2, 18, 34, 50, 66], [5, 1]),
np.reshape([0, 16, 32, 48, 64], [5, 1]),
np.reshape([self.init_val] * 5, [5, 1]),
np.reshape([1, 17, 33, 49, 65], [5, 1]),
np.reshape([self.init_val] * 5, [5, 1])
],
axis=1)
# The new weight matrix is of size
# [5-sized input layer, 4 class vocab + 1 class OOV].
remapped_matrix = variable_scope.get_variable(
name='dnn_output/obtained_weight_matrix',
shape=[5, 5],
initializer=loading_initializer,
partitioner=partitioned_variables.fixed_size_partitioner(2))
with self.test_session():
variables.global_variables_initializer().run()
self.assertAllClose(expected_remapped_matrix,
remapped_matrix.as_tensor().eval())
def test_load_and_remap_linear_multiclass_initializer_default_init(self):
"""Tests where the zeros_initializer default is used for linear."""
loading_initializer = (checkpoint_ops._load_and_remap_matrix_initializer(
new_row_vocab_size=5,
new_col_vocab_file=self.new_class_vocab_file,
old_col_vocab_file=self.old_class_vocab_file,
new_col_vocab_size=4,
old_tensor_name='some_scope/embeddings',
ckpt_path=[self.checkpoint_file],
new_row_vocab_file=self.new_feature_vocab_file,
old_row_vocab_file=self.old_feature_vocab_file,
num_row_oov_buckets=1,
num_col_oov_buckets=1))
expected_remapped_matrix = np.concatenate(
[
np.reshape([2, 18, 34, 50, 0, 0], [6, 1]),
np.reshape([0, 16, 32, 48, 0, 0], [6, 1]),
np.reshape([0] * 6, [6, 1]),
np.reshape([1, 17, 33, 49, 0, 0], [6, 1]),
np.reshape([0] * 6, [6, 1])
],
axis=1)
remapped_matrix = variable_scope.get_variable(
name='linear_init_fallback/obtained_weight_matrix',
shape=[6, 5],
initializer=loading_initializer,
partitioner=partitioned_variables.fixed_size_partitioner(2))
with self.test_session():
variables.global_variables_initializer().run()
self.assertAllClose(expected_remapped_matrix,
remapped_matrix.as_tensor().eval())
def test_load_embedding_initializer(self):
"""Tests for the load_embedding_initializer wrapper."""
embedding_loading_initializer = (checkpoint_ops._load_embedding_initializer(
new_vocab_file=self.new_feature_vocab_file,
old_vocab_file=self.old_feature_vocab_file,
new_vocab_size=5,
embedding_dim=16,
embedding_tensor_name='some_scope/embeddings',
ckpt_path=[self.checkpoint_file],
num_oov_buckets=1,
initializer=self.initializer))
expected_remapped_embeddings = np.concatenate(
[
np.reshape(range(64), [4, 16]),
np.reshape([self.init_val] * 32, [2, 16]),
],
axis=0)
# The new weight matrix is of size
# [5 feature vocab + 1 feature OOV, 16 (embedding dimension)], where the
# last vocab row (2nd last row) is newly initialized (wasn't found in
# previous vocab) and the actual last row is OOV and also newly initialized.
# Use a partitioned variable to confirm that the offset logic works.
remapped_embeddings = variable_scope.get_variable(
name='embedding/obtained_embedding_matrix',
shape=[6, 16],
initializer=embedding_loading_initializer,
partitioner=partitioned_variables.fixed_size_partitioner(2))
with self.test_session():
variables.global_variables_initializer().run()
self.assertAllClose(expected_remapped_embeddings,
remapped_embeddings.as_tensor().eval())
def testFixedSizePartitionerInt64(self):
with self.test_session():
partitioner = partitioned_variables.fixed_size_partitioner(4, axis=0)
with variable_scope.variable_scope("root", partitioner=partitioner):
v0 = variable_scope.get_variable("v0", dtype=dtypes.int64, shape=[20])
v0_list = v0._get_variable_list()
self.assertEqual(len(v0_list), 4)
def _testMultiplePartitionedVariables(self, is_training):
# When weights are partitioned into multiple partitions the weights variable
# is followed by a identity -> concat -> identity to group the partitions.
partitioner = partitioned_variables.fixed_size_partitioner(2)
graph = ops.Graph()
with graph.as_default():
with variable_scope.variable_scope('part', partitioner=partitioner):
batch_size, height, width, depth = 5, 128, 128, 3
input1 = array_ops.zeros((batch_size, height, width, depth))
input2 = array_ops.zeros((batch_size, height / 2, width / 2, 32))
conv = conv2d(
input1,
32, [5, 5],
stride=2,
padding='SAME',
weights_initializer=self._WeightInit(0.09),
activation_fn=None,
scope='test/test')
node = math_ops.add(conv, input2, name='test/add')
node = nn_ops.relu6(node, name='test/relu6')
quantize.Quantize(graph, is_training, weight_bits=8, activation_bits=8)
# Check that the weight's quant node was added.
op_names = [op.name for op in graph.get_operations()]
self.assertTrue(
'part/test/test/weights_quant/FakeQuantWithMinMaxVars' in op_names)
def test_loading_partitions_greater_than_max_rows(self):
"""Tests loading partitioned var with more slices than partitions."""
self._test_loading_variable_with_max_rows(
np_value=np.reshape(list(range(0, 36)), (9, 4)),
partitioner=partitioned_variables.fixed_size_partitioner(3),
# Even though each partition has 3 rows, we'll only load the tensor one
# row at a time.
max_rows_in_memory=1)
def test_loading_partitions_equals_max_rows(self):
"""Tests loading partitioned var sliced on partition boundary."""
self._test_loading_variable_with_max_rows(
np_value=np.reshape(list(range(0, 36)), (9, 4)),
partitioner=partitioned_variables.fixed_size_partitioner(3),
# With a tensor of shape [9, 3] and 3 partitions, each partition has
# exactly 3 rows.
max_rows_in_memory=3)
def test_loading_partitions_less_than_max_rows(self):
"""Tests loading partitioned var as a single slice.
(When the specified max_rows_in_memory is larger than the number of rows)
"""
self._test_loading_variable_with_max_rows(
np_value=np.reshape(list(range(0, 36)), (9, 4)),
partitioner=partitioned_variables.fixed_size_partitioner(3),
max_rows_in_memory=10)
def testFixedSizePartitioner(self):
with self.test_session():
partitioner = partitioned_variables.fixed_size_partitioner(5, axis=0)
with variable_scope.variable_scope("root", partitioner=partitioner):
v0 = variable_scope.get_variable(
"v0", dtype=dtypes.float32, shape=(10, 10))
v0_list = v0._get_variable_list()
v0_part = v0._get_partitions()
self.assertEqual(len(v0_list), 5)
self.assertAllEqual(v0_part, (5, 1))
def testFixedSizePartitionerInt64(self):
with self.cached_session():
partitioner = partitioned_variables.fixed_size_partitioner(4,
axis=0)
with variable_scope.variable_scope("root",
partitioner=partitioner):
v0 = variable_scope.get_variable("v0",
dtype=dtypes.int64,
shape=[20])
v0_list = v0._get_variable_list()
self.assertEqual(len(v0_list), 4)
def testResourceFixedSizePartitioner(self):
with self.cached_session():
partitioner = partitioned_variables.fixed_size_partitioner(5, axis=0)
with variable_scope.variable_scope(
"root", partitioner=partitioner, use_resource=True):
v0 = variable_scope.get_variable(
"v0", dtype=dtypes.float32, shape=(10, 10))
v0_list = v0._get_variable_list()
v0_part = v0._get_partitions()
self.assertEqual(len(v0_list), 5)
self.assertAllEqual(v0_part, (5, 1))
def test_dense_feature_with_partitioner(self):
with context.eager_mode():
sparse_input = sparse_tensor.SparseTensor(indices=((0, 0), (1, 0),
(2, 0), (3, 0)),
values=(0, 1, 3, 2),
dense_shape=(4, 4))
# Create feature columns (categorical and embedding).
categorical_column = fc.categorical_column_with_identity(
key='a', num_buckets=4)
embedding_dimension = 2
def _embedding_column_initializer(shape,
dtype,
partition_info=None):
offset = partition_info._var_offset[0]
del shape # unused
del dtype # unused
if offset == 0:
embedding_values = (
(1, 0), # id 0
(0, 1)) # id 1
else:
embedding_values = (
(1, 1), # id 2
(2, 2)) # id 3
return embedding_values
embedding_column = fc.embedding_column(
categorical_column,
dimension=embedding_dimension,
initializer=_embedding_column_initializer)
dense_features = df.DenseFeatures(
[embedding_column],
partitioner=partitioned_variables.fixed_size_partitioner(2))
features = {'a': sparse_input}
inputs = dense_features(features)
variables = dense_features.variables
# Sanity check: test that the inputs are correct.
self.assertAllEqual([[1, 0], [0, 1], [2, 2], [1, 1]], inputs)
# Check that only one variable was created.
self.assertEqual(2, len(variables))
# Check that invoking dense_features on the same features does not create
# additional variables
_ = dense_features(features)
self.assertEqual(2, len(variables))
self.assertIs(variables[0], dense_features.variables[0])
self.assertIs(variables[1], dense_features.variables[1])
def build(self, inputs_shape):
if inputs_shape[1].value is None:
raise ValueError("Expected inputs.shape[-1] to be known, saw shape: %s"
% inputs_shape)
input_depth = inputs_shape[1].value
h_depth = self._num_units if self._num_proj is None else self._num_proj
maybe_partitioner = (
partitioned_variables.fixed_size_partitioner(self._num_unit_shards)
if self._num_unit_shards is not None
else None)
self._kernel = self.add_variable(
_WEIGHTS_VARIABLE_NAME,
shape=[input_depth + h_depth, 4 * self._num_units],
initializer=self._initializer,
partitioner=maybe_partitioner)
self._bias = self.add_variable(
_BIAS_VARIABLE_NAME,
shape=[4 * self._num_units],
initializer=init_ops.zeros_initializer(dtype=self.dtype))
if self._use_peepholes:
self._w_f_diag = self.add_variable("w_f_diag", shape=[self._num_units],
initializer=self._initializer)
self._w_i_diag = self.add_variable("w_i_diag", shape=[self._num_units],
initializer=self._initializer)
self._w_o_diag = self.add_variable("w_o_diag", shape=[self._num_units],
initializer=self._initializer)
if self._num_proj is not None:
maybe_proj_partitioner = (
partitioned_variables.fixed_size_partitioner(self._num_proj_shards)
if self._num_proj_shards is not None
else None)
self._proj_kernel = self.add_variable(
"projection/%s" % _WEIGHTS_VARIABLE_NAME,
shape=[self._num_units, self._num_proj],
initializer=self._initializer,
partitioner=maybe_proj_partitioner)
self.built = True
def build(self, inputs_shape):
if inputs_shape[1].value is None:
raise ValueError("Expected inputs.shape[-1] to be known, saw shape: %s"
% inputs_shape)
input_depth = inputs_shape[1].value
h_depth = self._num_units if self._num_proj is None else self._num_proj
maybe_partitioner = (
partitioned_variables.fixed_size_partitioner(self._num_unit_shards)
if self._num_unit_shards is not None
else None)
self._kernel = self.add_variable(
_WEIGHTS_VARIABLE_NAME,
shape=[input_depth + h_depth, 4 * self._num_units],
initializer=self._initializer,
partitioner=maybe_partitioner)
self._bias = self.add_variable(
_BIAS_VARIABLE_NAME,
shape=[4 * self._num_units],
initializer=init_ops.zeros_initializer(dtype=self.dtype))
if self._use_peepholes:
self._w_f_diag = self.add_variable("w_f_diag", shape=[self._num_units],
initializer=self._initializer)
self._w_i_diag = self.add_variable("w_i_diag", shape=[self._num_units],
initializer=self._initializer)
self._w_o_diag = self.add_variable("w_o_diag", shape=[self._num_units],
initializer=self._initializer)
if self._num_proj is not None:
maybe_proj_partitioner = (
partitioned_variables.fixed_size_partitioner(self._num_proj_shards)
if self._num_proj_shards is not None
else None)
self._proj_kernel = self.add_variable(
"projection/%s" % _WEIGHTS_VARIABLE_NAME,
shape=[self._num_units, self._num_proj],
initializer=self._initializer,
partitioner=maybe_proj_partitioner)
self.built = True
def make_graph_with_partitioned_variables(use_resource):
variable_scope.get_variable(
name="weights",
partitioner=partitioned_variables.fixed_size_partitioner(3, axis=0),
initializer=random_ops.truncated_normal([100, 10]),
use_resource=use_resource)
# The next variable illustrates the necessity of restoring collections
# in a deterministic fashion when using ResourceVariables.
variable_scope.get_variable(
name="another",
shape=[],
collections=["a", "b", "z", "f", "e", "d", "g"],
use_resource=use_resource)
def make_graph_with_partitioned_variables(use_resource):
variable_scope.get_variable(
name="weights",
partitioner=partitioned_variables.fixed_size_partitioner(3, axis=0),
initializer=random_ops.truncated_normal([100, 10]),
use_resource=use_resource)
# The next variable illustrates the necessity of restoring collections
# in a deterministic fashion when using ResourceVariables.
variable_scope.get_variable(
name="another",
shape=[],
collections=["a", "b", "z", "f", "e", "d", "g"],
use_resource=use_resource)
def _random_weights(self, size=50, num_shards=1):
assert size > 0
assert num_shards > 0
assert num_shards <= size
embedding_weights = list(variable_scope.get_variable(
"embedding_weights",
shape=[size],
partitioner=partitioned_variables.fixed_size_partitioner(num_shards),
initializer=init_ops.truncated_normal_initializer(
mean=0.0, stddev=1.0, dtype=dtypes.float32)))
for w in embedding_weights:
w.initializer.run()
return embedding_weights
def testVariablesNotParitioned_MovingAvg(self):
# Variables added should not use a default partiioner since they are
# scalar. There would be a tensorflow error thrown if the partitioner was
# respected by the rewrite.
with ops.Graph().as_default():
with variable_scope.variable_scope(
'part', partitioner=partitioned_variables.fixed_size_partitioner(2)):
x = array_ops.placeholder(dtypes.float32, shape=[2])
_ = quant_ops.MovingAvgQuantize(
x,
init_min=0.0,
init_max=0.0,
is_training=True,
vars_collection=_MIN_MAX_VARS)
def testVariablesNotPartitioned_LastValue(self):
# Variables added should not use a default partiioner since they are
# scalar. There would be a tensorflow error thrown if the partitioner was
# respected by the rewrite.
with ops.Graph().as_default():
with variable_scope.variable_scope(
'part', partitioner=partitioned_variables.fixed_size_partitioner(2)):
x = array_ops.placeholder(dtypes.float32, shape=[2])
_ = quant_ops.LastValueQuantize(
x,
init_min=0.0,
init_max=0.0,
is_training=True,
vars_collection=_MIN_MAX_VARS)
def testReusePartitionedVaraiblesAndRegularizers(self):
regularizer = lambda x: math_ops.reduce_sum(x) * 1e-3
partitioner = partitioned_variables.fixed_size_partitioner(3)
for reuse in [False, True]:
with variable_scope.variable_scope(variable_scope.get_variable_scope(),
partitioner=partitioner,
reuse=reuse):
layer = base_layers.Layer(name='my_layer')
_ = layer.add_variable(
'reg_part_var', [4, 4],
initializer=init_ops.zeros_initializer(),
regularizer=regularizer)
self.assertEqual(
len(ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)), 3)
def testNumPartitionsLargerThanSize(self):
strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
self.cluster_resolver, partitioned_variables.fixed_size_partitioner(4))
with strategy.scope():
v = variables.Variable([0, 1, 2])
self.assertIsInstance(v, sharded_variable.ShardedVariable)
self.assertLen(v.variables, 3)
self.assertRegex(v.variables[0].device, "/job:ps/replica:0/task:0")
self.assertRegex(v.variables[1].device, "/job:ps/replica:0/task:1")
self.assertRegex(v.variables[2].device, "/job:ps/replica:0/task:0")
self.assertAllEqual(v.variables[0].read_value().numpy(), [0])
self.assertAllEqual(v.variables[1].read_value().numpy(), [1])
self.assertAllEqual(v.variables[2].read_value().numpy(), [2])
def testReusePartitionedVaraiblesAndRegularizers(self):
regularizer = lambda x: math_ops.reduce_sum(x) * 1e-3
partitioner = partitioned_variables.fixed_size_partitioner(3)
for reuse in [False, True]:
with variable_scope.variable_scope(variable_scope.get_variable_scope(),
partitioner=partitioner,
reuse=reuse):
layer = base_layers.Layer(name='my_layer')
variable = layer.add_variable(
'reg_part_var', [4, 4],
initializer=init_ops.zeros_initializer(),
regularizer=regularizer)
self.assertEqual(
len(ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)), 3)
def testColocateWith(self):
strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
self.cluster_resolver, partitioned_variables.fixed_size_partitioner(2))
with strategy.scope():
v1 = variables.Variable([0, 1, 2, 3])
with strategy.extended.colocate_vars_with(v1.variables[0]):
v2 = variables.Variable([4, 5])
self.assertIsInstance(v1, sharded_variable.ShardedVariable)
self.assertIsInstance(v2, variables.Variable)
self.assertNotIsInstance(v2, sharded_variable.ShardedVariable)
self.assertEqual(v2.device, v1.variables[0].device)
self.assertAllEqual(v2.read_value().numpy(), [4, 5])
def _ShardTestEmbeddings(self, weights, biases, num_shards):
"""Shards the weights and biases returned by _GenerateTestData.
Args:
weights: The weights returned by _GenerateTestData.
biases: The biases returned by _GenerateTestData.
num_shards: The number of shards to create.
Returns:
sharded_weights: A list of size `num_shards` containing all the weights.
sharded_biases: A list of size `num_shards` containing all the biases.
"""
with ops.Graph().as_default() as g:
sharded_weights = variable_scope.get_variable(
"w",
partitioner=partitioned_variables.fixed_size_partitioner(num_shards),
initializer=constant_op.constant(weights))
sharded_biases = variable_scope.get_variable(
"b",
partitioner=partitioned_variables.fixed_size_partitioner(num_shards),
initializer=constant_op.constant(biases))
with self.test_session(graph=g) as sess:
variables.global_variables_initializer().run()
return sess.run([list(sharded_weights), list(sharded_biases)])
def _test_device_assignment_distributed_enable_partitioner(
self, task_type, task_id, num_gpus, use_core_strategy=False):
d, _, sess_config = self._get_test_objects(
task_type, task_id, num_gpus, use_core_strategy=use_core_strategy)
num_shards = len(d.extended.parameter_devices)
partitioner = partitioned_variables.fixed_size_partitioner(num_shards)
with ops.Graph().as_default(), \
self.cached_session(target=self._default_target,
config=sess_config) as sess, \
d.scope():
n = variable_scope.get_variable(
'n',
initializer=constant_op.constant([10.0, 20.0]),
aggregation=variable_scope.VariableAggregation.SUM,
partitioner=partitioner)
for part_id, var in enumerate(n):
self.assertEqual(var.device, '/job:ps/task:%d' % part_id)
def model_fn():
a = constant_op.constant([3.0, 5.0])
# The device scope is ignored for variables but not for normal ops.
with ops.device('/job:worker/task:0'):
x = variable_scope.get_variable(
'x',
initializer=constant_op.constant([10.0, 20.0]),
aggregation=variable_scope.VariableAggregation.SUM,
partitioner=partitioner)
x_add = x.assign_add(a, name='x_add')
# The variable x is on the task 1 since the device_function has been
# called once before the model_fn.
for part_id, var in enumerate(x):
self.assertEqual(var.device, '/job:ps/task:%d' % part_id)
self.assertEqual(var.device, x_add[part_id].device)
return x_add
x = d.extended.call_for_each_replica(model_fn)
if context.num_gpus() >= 1:
variables.global_variables_initializer().run()
x_val = sess.run(x)
if num_gpus < 1:
self.assertEqual(x_val, [13.0, 25.0])
else:
x_expect = [10.0 + 3 * num_gpus, 20.0 + 5 * num_gpus]
self.assertEqual(x_val, x_expect)
def _test_device_assignment_distributed_enable_partitioner(
self, task_type, task_id, num_gpus, use_core_strategy=False):
d, _, sess_config = self._get_test_objects(
task_type, task_id, num_gpus, use_core_strategy=use_core_strategy)
num_shards = len(d.extended.parameter_devices)
partitioner = partitioned_variables.fixed_size_partitioner(num_shards)
with ops.Graph().as_default(), \
self.cached_session(target=self._default_target,
config=sess_config) as sess, \
d.scope():
n = variable_scope.get_variable(
'n',
initializer=constant_op.constant([10.0, 20.0]),
aggregation=variable_scope.VariableAggregation.SUM,
partitioner=partitioner)
for part_id, var in enumerate(n):
self.assertEqual(var.device, '/job:ps/task:%d' % part_id)
def model_fn():
a = constant_op.constant([3.0, 5.0])
# The device scope is ignored for variables but not for normal ops.
with ops.device('/job:worker/task:0'):
x = variable_scope.get_variable(
'x',
initializer=constant_op.constant([10.0, 20.0]),
aggregation=variable_scope.VariableAggregation.SUM,
partitioner=partitioner)
x_add = x.assign_add(a, name='x_add')
# The variable x is on the task 1 since the device_function has been
# called once before the model_fn.
for part_id, var in enumerate(x):
self.assertEqual(var.device, '/job:ps/task:%d' % part_id)
self.assertEqual(var.device, x_add[part_id].device)
return x_add
x = d.extended.call_for_each_replica(model_fn)
if context.num_gpus() >= 1:
variables.global_variables_initializer().run()
x_val = sess.run(x)
if num_gpus < 1:
self.assertEqual(x_val, [13.0, 25.0])
else:
x_expect = [10.0 + 3 * num_gpus, 20.0 + 5 * num_gpus]
self.assertEqual(x_val, x_expect)
def testVariablesNotPartitioned_MovingAvg(self):
# Variables added should not use a default partiioner since they are
# scalar. There would be a tensorflow error thrown if the partitioner was
# respected by the rewrite.
with ops.Graph().as_default():
with variable_scope.variable_scope(
'part',
partitioner=partitioned_variables.fixed_size_partitioner(
2)):
x = array_ops.placeholder(dtypes.float32, shape=[2])
min_var = variables.Variable(0.0)
max_var = variables.Variable(0.0)
_ = quant_ops.MovingAvgQuantize(x,
min_var,
max_var,
is_training=True)
def testNonCallableInitialValue(self):
strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
self.cluster_resolver, partitioned_variables.fixed_size_partitioner(4))
with strategy.scope():
v = variables.Variable([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
self.assertIsInstance(v, sharded_variable.ShardedVariable)
self.assertLen(v.variables, 4)
self.assertRegex(v.variables[0].device, "/job:ps/replica:0/task:0")
self.assertRegex(v.variables[1].device, "/job:ps/replica:0/task:1")
self.assertRegex(v.variables[2].device, "/job:ps/replica:0/task:0")
self.assertRegex(v.variables[3].device, "/job:ps/replica:0/task:1")
self.assertAllEqual(v.variables[0].read_value().numpy(), [0, 1, 2])
self.assertAllEqual(v.variables[1].read_value().numpy(), [3, 4, 5])
self.assertAllEqual(v.variables[2].read_value().numpy(), [6, 7])
self.assertAllEqual(v.variables[3].read_value().numpy(), [8, 9])
def testCreateSlotFromScalarVariable(self):
with self.test_session():
s = variables.Variable(1.0, name="var")
p_v = variable_scope.get_variable(
"var",
shape=[2, 2],
partitioner=partitioned_variables.fixed_size_partitioner(2))
for i, v in enumerate(p_v):
slot = slot_creator.create_slot(v, s.initialized_value(), name="slot")
variables.global_variables_initializer().run()
self.assertEqual("var/part_%d/slot" % i, slot.op.name)
self.assertEqual([], slot.get_shape().as_list())
self.assertEqual(dtypes.float32, slot.dtype.base_dtype)
self.assertAllEqual(1.0, slot.eval())
def _random_weights(self, size=50, num_shards=1):
assert size > 0
assert num_shards > 0
assert num_shards <= size
embedding_weights = list(
variable_scope.get_variable(
"embedding_weights",
shape=[size],
partitioner=partitioned_variables.fixed_size_partitioner(
num_shards),
initializer=init_ops.truncated_normal_initializer(
mean=0.0, stddev=1.0, dtype=dtypes.float32)))
for w in embedding_weights:
w.initializer.run()
return embedding_weights
def _random_weights(self, vocab_size=4, embed_dim=4, num_shards=1):
assert vocab_size > 0
assert embed_dim > 0
assert num_shards > 0
assert num_shards <= vocab_size
initializer = init_ops.truncated_normal_initializer(
mean=0.0, stddev=1.0 / math.sqrt(vocab_size), dtype=dtypes.float32)
embedding_weights = list(variable_scope.get_variable(
name="embedding_weights",
shape=[vocab_size, embed_dim],
partitioner=partitioned_variables.fixed_size_partitioner(num_shards),
initializer=initializer))
for w in embedding_weights:
self.evaluate(w.initializer)
embedding_weights = [self.evaluate(w) for w in embedding_weights]
return embedding_weights
def test_load_embedding_initializer_large_oov(self):
"""Tests for the large OOV case for load_embedding_initializer wrapper."""
self.new_feature_vocab_file = os.path.join(self.get_temp_dir(),
'new_feature_vocab.txt')
with open(self.new_feature_vocab_file, 'w') as f:
f.write('\n'.join(['one', 'zero', 'two', 'four']) + '\n')
# Checkpoint has 5 entries, 3 of which correspond to OOV.
self.old_feature_vocab_file = os.path.join(self.get_temp_dir(),
'old_feature_vocab.txt')
with open(self.old_feature_vocab_file, 'w') as f:
f.write('\n'.join(['zero', 'one']) + '\n')
embedding_loading_initializer = (
checkpoint_ops._load_embedding_initializer(
new_vocab_file=self.new_feature_vocab_file,
old_vocab_file=self.old_feature_vocab_file,
new_vocab_size=4,
embedding_dim=16,
embedding_tensor_name='some_scope/embeddings',
ckpt_path=[self.checkpoint_file],
num_oov_buckets=5,
initializer=self.initializer))
expected_remapped_embeddings = np.concatenate([
np.reshape(range(16, 32), [1, 16]),
np.reshape(range(16), [1, 16]),
np.reshape([self.init_val] * 112, [7, 16]),
],
axis=0)
# The new weight matrix is of size
# [4 feature vocab + 5 feature OOV, 16 (embedding dimension)], where the
# 3rd and 4th rows are not found in the old vocabulary and therefore newly
# initialized. The last five rows are OOV and also newly initialized.
# Use a partitioned variable to confirm that the offset logic works.
remapped_embeddings = variable_scope.get_variable(
name='embedding/obtained_embedding_matrix',
shape=[9, 16],
initializer=embedding_loading_initializer,
partitioner=partitioned_variables.fixed_size_partitioner(2))
with self.test_session():
variables.global_variables_initializer().run()
self.assertAllClose(expected_remapped_embeddings,
remapped_embeddings.as_tensor().eval())
def test_load_embedding_initializer_large_oov(self):
"""Tests for the large OOV case for load_embedding_initializer wrapper."""
self.new_feature_vocab_file = os.path.join(
self.get_temp_dir(), 'new_feature_vocab.txt')
with open(self.new_feature_vocab_file, 'w') as f:
f.write('\n'.join(['one', 'zero', 'two', 'four']) + '\n')
# Checkpoint has 5 entries, 3 of which correspond to OOV.
self.old_feature_vocab_file = os.path.join(
self.get_temp_dir(), 'old_feature_vocab.txt')
with open(self.old_feature_vocab_file, 'w') as f:
f.write('\n'.join(['zero', 'one']) + '\n')
embedding_loading_initializer = (checkpoint_ops._load_embedding_initializer(
new_vocab_file=self.new_feature_vocab_file,
old_vocab_file=self.old_feature_vocab_file,
new_vocab_size=4,
embedding_dim=16,
embedding_tensor_name='some_scope/embeddings',
ckpt_path=[self.checkpoint_file],
num_oov_buckets=5,
initializer=self.initializer))
expected_remapped_embeddings = np.concatenate(
[
np.reshape(range(16, 32), [1, 16]),
np.reshape(range(16), [1, 16]),
np.reshape([self.init_val] * 112, [7, 16]),
],
axis=0)
# The new weight matrix is of size
# [4 feature vocab + 5 feature OOV, 16 (embedding dimension)], where the
# 3rd and 4th rows are not found in the old vocabulary and therefore newly
# initialized. The last five rows are OOV and also newly initialized.
# Use a partitioned variable to confirm that the offset logic works.
remapped_embeddings = variable_scope.get_variable(
name='embedding/obtained_embedding_matrix',
shape=[9, 16],
initializer=embedding_loading_initializer,
partitioner=partitioned_variables.fixed_size_partitioner(2))
with self.test_session():
variables.global_variables_initializer().run()
self.assertAllClose(expected_remapped_embeddings,
remapped_embeddings.as_tensor().eval())
def _random_weights(self, vocab_size=4, embed_dim=4, num_shards=1):
assert vocab_size > 0
assert embed_dim > 0
assert num_shards > 0
assert num_shards <= vocab_size
initializer = init_ops.truncated_normal_initializer(
mean=0.0, stddev=1.0 / math.sqrt(vocab_size), dtype=dtypes.float32)
embedding_weights = list(variable_scope.get_variable(
"embedding_weights",
shape=[vocab_size, embed_dim],
partitioner=partitioned_variables.fixed_size_partitioner(num_shards),
initializer=initializer))
for w in embedding_weights:
w.initializer.run()
embedding_weights = [w.eval() for w in embedding_weights]
return embedding_weights
def testPartitionAwareInitializer(self):
strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
self.cluster_resolver, partitioned_variables.fixed_size_partitioner(2))
with strategy.scope():
initializer = PartitionAwareIdentity()
initial_value = functools.partial(
initializer, shape=(4, 4), dtype=dtypes.int64)
v = variables.Variable(
initial_value=initial_value, shape=(4, 4), dtype=dtypes.int64)
self.assertIsInstance(v, sharded_variable.ShardedVariable)
self.assertLen(v.variables, 2)
self.assertRegex(v.variables[0].device, "/job:ps/replica:0/task:0")
self.assertRegex(v.variables[1].device, "/job:ps/replica:0/task:1")
self.assertAllEqual(v.variables[0].read_value().numpy(),
[[1, 0, 0, 0], [0, 1, 0, 0]])
self.assertAllEqual(v.variables[1].read_value().numpy(),
[[0, 0, 1, 0], [0, 0, 0, 1]])
def make_variable_dict(max_age, max_gender, partitioned=False):
# TODO(sibyl-toe9oF2e): Figure out how to derive max_age & max_gender from
# examples_dict.
partitioner = None
if partitioned:
partitioner = partitioned_variables.fixed_size_partitioner(num_shards=2,
axis=0)
with variable_scope.variable_scope(
name_or_scope='variables',
partitioner=partitioner):
age_weights = variables_lib.Variable(
array_ops.zeros(
[max_age + 1], dtype=dtypes.float32))
gender_weights = variables_lib.Variable(
array_ops.zeros(
[max_gender + 1], dtype=dtypes.float32))
return dict(
sparse_features_weights=[age_weights, gender_weights],
dense_features_weights=[])
def testLargePartitionedVariables(self):
save_path = os.path.join(self.get_temp_dir(), "large_variable")
var_name = "my_var"
# Saving large partition variable.
with session.Session("", graph=ops.Graph()) as sess:
with ops.device("/cpu:0"):
# Create a partitioned variable which is larger than int32 size but
# split into smaller sized variables.
init = lambda shape, dtype, partition_info: constant_op.constant(
True, dtype, shape)
partitioned_var = list(variable_scope.get_variable(
var_name,
shape=[1 << 31],
partitioner=partitioned_variables.fixed_size_partitioner(4),
initializer=init,
dtype=dtypes.bool))
variables.global_variables_initializer().run()
save = saver.Saver(partitioned_var)
val = save.save(sess, save_path)
self.assertEqual(save_path, val)
def testCreateSlotFromFirstMDimensionVariable(self):
with self.test_session():
s = variables.Variable([1.0, 2.5], name="var")
p_v = variable_scope.get_variable(
"var",
shape=[2, 2],
partitioner=partitioned_variables.fixed_size_partitioner(2))
for i, v in enumerate(p_v):
slot = slot_creator.create_slot(v, s.initialized_value(), name="slot")
si = slot._save_slice_info
variables.global_variables_initializer().run()
self.assertEqual("var/part_%d/slot" % i, slot.op.name)
self.assertEqual([2], slot.get_shape().as_list())
self.assertEqual(dtypes.float32, slot.dtype.base_dtype)
self.assertAllEqual([1.0, 2.5], slot.eval())
self.assertAllEqual([2], si.full_shape)
self.assertAllEqual([i], si.var_offset)
self.assertAllEqual([1], si.var_shape)
def testMixedFeaturesArbitraryWeightsPartitioned(self):
"""Tests SDCALinearRegressor works with a mix of features (partitioned)."""
def input_fn():
return {
'example_id':
constant_op.constant(['1', '2', '3']),
'price':
constant_op.constant([[0.6], [0.8], [0.3]]),
'sq_footage':
constant_op.constant([[900.0], [700.0], [600.0]]),
'country':
sparse_tensor.SparseTensor(
values=['IT', 'US', 'GB'],
indices=[[0, 0], [1, 3], [2, 1]],
dense_shape=[3, 5]),
'weights':
constant_op.constant([[3.0], [5.0], [7.0]])
}, constant_op.constant([[1.55], [-1.25], [-3.0]])
with self._single_threaded_test_session():
price = feature_column_lib.real_valued_column('price')
sq_footage_bucket = feature_column_lib.bucketized_column(
feature_column_lib.real_valued_column('sq_footage'),
boundaries=[650.0, 800.0])
country = feature_column_lib.sparse_column_with_hash_bucket(
'country', hash_bucket_size=5)
sq_footage_country = feature_column_lib.crossed_column(
[sq_footage_bucket, country], hash_bucket_size=10)
regressor = sdca_estimator.SDCALinearRegressor(
example_id_column='example_id',
feature_columns=[
price, sq_footage_bucket, country, sq_footage_country
],
l2_regularization=1.0,
weight_column_name='weights',
partitioner=partitioned_variables.fixed_size_partitioner(
num_shards=2, axis=0))
regressor.fit(input_fn=input_fn, steps=20)
loss = regressor.evaluate(input_fn=input_fn, steps=1)['loss']
self.assertLess(loss, 0.05)
def testPartitionedVariables(self):
"""Tests LinearClassifier with LinearSDCA with partitioned variables."""
def input_fn():
return {
'example_id':
constant_op.constant(['1', '2', '3']),
'price':
constant_op.constant([[0.6], [0.8], [0.3]]),
'sq_footage':
constant_op.constant([[900.0], [700.0], [600.0]]),
'country':
sparse_tensor.SparseTensor(
values=['IT', 'US', 'GB'],
indices=[[0, 0], [1, 3], [2, 1]],
dense_shape=[3, 5]),
'weights':
constant_op.constant([[3.0], [1.0], [1.0]])
}, constant_op.constant([[1], [0], [1]])
price = feature_column_v2.numeric_column_v2('price')
sq_footage_bucket = feature_column_v2.bucketized_column_v2(
feature_column_v2.numeric_column_v2('sq_footage'),
boundaries=[650.0, 800.0])
country = feature_column_v2.categorical_column_with_hash_bucket_v2(
'country', hash_bucket_size=5)
sq_footage_country = feature_column_v2.crossed_column_v2(
[sq_footage_bucket, 'country'], hash_bucket_size=10)
optimizer = linear.LinearSDCA(
example_id_column='example_id', symmetric_l2_regularization=0.01)
classifier = linear.LinearClassifier(
feature_columns=[price, sq_footage_bucket, country, sq_footage_country],
weight_column='weights',
partitioner=partitioned_variables.fixed_size_partitioner(
num_shards=2, axis=0),
optimizer=optimizer)
classifier.train(input_fn=input_fn, steps=100)
loss = classifier.evaluate(input_fn=input_fn, steps=1)['loss']
self.assertLess(loss, 0.2)
def testMixedFeaturesArbitraryWeightsPartitioned(self):
"""Tests SDCALinearRegressor works with a mix of features (partitioned)."""
def input_fn():
return {
'example_id':
constant_op.constant(['1', '2', '3']),
'price':
constant_op.constant([[0.6], [0.8], [0.3]]),
'sq_footage':
constant_op.constant([[900.0], [700.0], [600.0]]),
'country':
sparse_tensor.SparseTensor(
values=['IT', 'US', 'GB'],
indices=[[0, 0], [1, 3], [2, 1]],
dense_shape=[3, 5]),
'weights':
constant_op.constant([[3.0], [5.0], [7.0]])
}, constant_op.constant([[1.55], [-1.25], [-3.0]])
with self._single_threaded_test_session():
price = feature_column_lib.real_valued_column('price')
sq_footage_bucket = feature_column_lib.bucketized_column(
feature_column_lib.real_valued_column('sq_footage'),
boundaries=[650.0, 800.0])
country = feature_column_lib.sparse_column_with_hash_bucket(
'country', hash_bucket_size=5)
sq_footage_country = feature_column_lib.crossed_column(
[sq_footage_bucket, country], hash_bucket_size=10)
regressor = sdca_estimator.SDCALinearRegressor(
example_id_column='example_id',
feature_columns=[
price, sq_footage_bucket, country, sq_footage_country
],
l2_regularization=1.0,
weight_column_name='weights',
partitioner=partitioned_variables.fixed_size_partitioner(
num_shards=2, axis=0))
regressor.fit(input_fn=input_fn, steps=20)
loss = regressor.evaluate(input_fn=input_fn, steps=1)['loss']
self.assertLess(loss, 0.05)
def _create_partitioned_variables(name,
num_hosts,
vocabulary_size,
embedding_dimension,
initializer,
collections=None): # pylint: disable=redefined-outer-name
"""Creates ParitionedVariables based on `num_hosts` for `table`."""
# TODO(shizhiw): automatically place embedding lookup elsewhere?
if vocabulary_size < num_hosts:
raise ValueError('`vocabulary_size`({}) is smaller than `num_hosts`({}). '
'As TPU embedding is not optimized for small tables, '
'please consider other ways for this embedding lookup.')
return list(variable_scope.get_variable(
name,
shape=(vocabulary_size, embedding_dimension),
partitioner=partitioned_variables.fixed_size_partitioner(num_hosts),
dtype=dtypes.float32,
initializer=initializer,
collections=collections,
trainable=False))
def _create_partitioned_variables(name,
num_hosts,
vocabulary_size,
embedding_dimension,
initializer,
collections=None): # pylint: disable=redefined-outer-name
"""Creates ParitionedVariables based on `num_hosts` for `table`."""
# TODO(shizhiw): automatically place embedding lookup elsewhere?
if vocabulary_size < num_hosts:
raise ValueError('`vocabulary_size`({}) is smaller than `num_hosts`({}). '
'As TPU embedding is not optimized for small tables, '
'please consider other ways for this embedding lookup.')
return list(variable_scope.get_variable(
name,
shape=(vocabulary_size, embedding_dimension),
partitioner=partitioned_variables.fixed_size_partitioner(num_hosts),
dtype=dtypes.float32,
initializer=initializer,
collections=collections,
trainable=False))
def testPartitionedMixedFeatures(self):
"""Tests SDCALogisticClassifier with a mix of features (partitioned)."""
def input_fn():
return {
'example_id':
constant_op.constant(['1', '2', '3']),
'price':
constant_op.constant([[0.6], [0.8], [0.3]]),
'sq_footage':
constant_op.constant([900.0, 700.0, 600.0]),
'country':
sparse_tensor.SparseTensor(
values=['IT', 'US', 'GB'],
indices=[[0, 0], [1, 3], [2, 1]],
dense_shape=[3, 5]),
'weights':
constant_op.constant([[3.0], [1.0], [1.0]])
}, constant_op.constant([[1], [0], [1]])
with self._single_threaded_test_session():
price = feature_column_lib.real_valued_column('price')
sq_footage_bucket = feature_column_lib.bucketized_column(
feature_column_lib.real_valued_column('sq_footage'),
boundaries=[650.0, 800.0])
country = feature_column_lib.sparse_column_with_hash_bucket(
'country', hash_bucket_size=5)
sq_footage_country = feature_column_lib.crossed_column(
[sq_footage_bucket, country], hash_bucket_size=10)
classifier = sdca_estimator.SDCALogisticClassifier(
example_id_column='example_id',
feature_columns=[
price, sq_footage_bucket, country, sq_footage_country
],
weight_column_name='weights',
partitioner=partitioned_variables.fixed_size_partitioner(
num_shards=2, axis=0))
classifier.fit(input_fn=input_fn, steps=50)
metrics = classifier.evaluate(input_fn=input_fn, steps=1)
self.assertGreater(metrics['accuracy'], 0.9)
def testPartitionedMixedFeatures(self):
"""Tests SDCALogisticClassifier with a mix of features (partitioned)."""
def input_fn():
return {
'example_id':
constant_op.constant(['1', '2', '3']),
'price':
constant_op.constant([[0.6], [0.8], [0.3]]),
'sq_footage':
constant_op.constant([900.0, 700.0, 600.0]),
'country':
sparse_tensor.SparseTensor(
values=['IT', 'US', 'GB'],
indices=[[0, 0], [1, 3], [2, 1]],
dense_shape=[3, 5]),
'weights':
constant_op.constant([[3.0], [1.0], [1.0]])
}, constant_op.constant([[1], [0], [1]])
with self._single_threaded_test_session():
price = feature_column_lib.real_valued_column('price')
sq_footage_bucket = feature_column_lib.bucketized_column(
feature_column_lib.real_valued_column('sq_footage'),
boundaries=[650.0, 800.0])
country = feature_column_lib.sparse_column_with_hash_bucket(
'country', hash_bucket_size=5)
sq_footage_country = feature_column_lib.crossed_column(
[sq_footage_bucket, country], hash_bucket_size=10)
classifier = sdca_estimator.SDCALogisticClassifier(
example_id_column='example_id',
feature_columns=[
price, sq_footage_bucket, country, sq_footage_country
],
weight_column_name='weights',
partitioner=partitioned_variables.fixed_size_partitioner(
num_shards=2, axis=0))
classifier.fit(input_fn=input_fn, steps=50)
metrics = classifier.evaluate(input_fn=input_fn, steps=1)
self.assertGreater(metrics['accuracy'], 0.9)
def test_initializer_with_oov_only_partition(self):
"""Tests for the output layer initializer where one partition is all OOV."""
loading_initializer = (checkpoint_ops._load_and_remap_matrix_initializer(
new_row_vocab_size=5,
new_col_vocab_file=self.new_class_vocab_file,
old_col_vocab_file=self.old_class_vocab_file,
new_col_vocab_size=4,
old_tensor_name='some_scope/embeddings',
ckpt_path=[self.checkpoint_file],
new_row_vocab_file=self.new_feature_vocab_file,
old_row_vocab_file=self.old_feature_vocab_file,
num_row_oov_buckets=5,
num_col_oov_buckets=1,
initializer=self.initializer))
expected_remapped_matrix = np.concatenate(
[
np.reshape([2, 18, 34, 50] + [self.init_val] * 6, [10, 1]),
np.reshape([0, 16, 32, 48] + [self.init_val] * 6, [10, 1]),
np.reshape([self.init_val] * 10, [10, 1]),
np.reshape([1, 17, 33, 49] + [self.init_val] * 6, [10, 1]),
np.reshape([self.init_val] * 10, [10, 1]),
],
axis=1)
# The new weight matrix is of size
# [5 feature vocab + 5 feature OOV, 4 class vocab + 1 class OOV]. The
# second partition has only OOV.
remapped_matrix = variable_scope.get_variable(
name='linear_all_oov/obtained_weight_matrix',
shape=[10, 5],
initializer=loading_initializer,
partitioner=partitioned_variables.fixed_size_partitioner(2))
with self.test_session():
variables.global_variables_initializer().run()
self.assertAllClose(expected_remapped_matrix,
remapped_matrix.as_tensor().eval())
def test_load_and_remap_output_layer_weight_initializer_linear(self):
"""Tests for the output layer initializer in the linear multi-class case."""
loading_initializer = (checkpoint_ops._load_and_remap_matrix_initializer(
new_row_vocab_size=5,
new_col_vocab_file=self.new_class_vocab_file,
old_col_vocab_file=self.old_class_vocab_file,
new_col_vocab_size=4,
old_tensor_name='some_scope/embeddings',
ckpt_path=[self.checkpoint_file],
new_row_vocab_file=self.new_feature_vocab_file,
old_row_vocab_file=self.old_feature_vocab_file,
num_row_oov_buckets=1,
num_col_oov_buckets=1,
initializer=self.initializer))
expected_remapped_matrix = np.concatenate(
[
np.reshape([2, 18, 34, 50, self.init_val, self.init_val], [6, 1]),
np.reshape([0, 16, 32, 48, self.init_val, self.init_val], [6, 1]),
np.reshape([self.init_val] * 6, [6, 1]),
np.reshape([1, 17, 33, 49, self.init_val, self.init_val], [6, 1]),
np.reshape([self.init_val] * 6, [6, 1])
],
axis=1)
# The new weight matrix is of size
# [5 feature vocab + 1 feature OOV, 4 class vocab + 1 class OOV]. Use a
# partitioned variable to confirm that the offset logic works.
remapped_matrix = variable_scope.get_variable(
name='linear/obtained_weight_matrix',
shape=[6, 5],
initializer=loading_initializer,
partitioner=partitioned_variables.fixed_size_partitioner(2))
with self.test_session():
variables.global_variables_initializer().run()
self.assertAllClose(expected_remapped_matrix,
remapped_matrix.as_tensor().eval())
def test_load_embedding_initializer_old_row_vocab(self):
"""Tests for load_embedding_initializer where we constrain old vocab."""
embedding_loading_initializer = (
checkpoint_ops._load_embedding_initializer(
new_vocab_file=self.new_feature_vocab_file,
old_vocab_file=self.old_feature_vocab_file,
# Considered old vocabulary becomes ['zero', 'one', 'two']. This
# means 'three' in the new vocabulary is newly initialized.
old_vocab_size=3,
new_vocab_size=5,
embedding_dim=16,
embedding_tensor_name='some_scope/embeddings',
ckpt_path=[self.checkpoint_file],
num_oov_buckets=1,
initializer=self.initializer))
# The new weight matrix is of size
# [5 feature vocab + 1 feature OOV, 16 (embedding dimension)], where the
# last vocab row (2nd last row) is newly initialized (wasn't found in
# previous vocab) and the actual last row is OOV and also newly initialized.
# Use a partitioned variable to confirm that the offset logic works.
expected_remapped_embeddings = np.concatenate(
[
np.reshape(range(48), [3, 16]),
np.reshape([self.init_val] * 48, [3, 16]),
],
axis=0)
remapped_embeddings = variable_scope.get_variable(
name='embedding/obtained_embedding_matrix',
shape=[6, 16],
initializer=embedding_loading_initializer,
partitioner=partitioned_variables.fixed_size_partitioner(2))
with self.cached_session():
self.evaluate(variables.global_variables_initializer())
self.assertAllClose(expected_remapped_embeddings,
remapped_embeddings.as_tensor())
def testPartitionWhenLackOfInfo(self):
strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
self.cluster_resolver, partitioned_variables.fixed_size_partitioner(2))
with strategy.scope():
initializer = init_ops_v2.Constant([0, 1, 2, 3])
# Shape is not explicitly specified.
v1 = variables.Variable(
initial_value=lambda: initializer(shape=(4,), dtype=dtypes.int64),
dtype=dtypes.int64)
# Dtype is not explicitly specified.
v2 = variables.Variable(
initial_value=lambda: initializer(shape=(4,), dtype=dtypes.int64),
shape=(4,))
# Neither shape nor dtype is explicitly specified.
v3 = variables.Variable(
initial_value=lambda: initializer(shape=(4,), dtype=dtypes.int64))
for v in [v1, v2, v3]:
self.assertIsInstance(v, sharded_variable.ShardedVariable)
self.assertLen(v.variables, 2)
self.assertRegex(v.variables[0].device, "/job:ps/replica:0/task:0")
self.assertRegex(v.variables[1].device, "/job:ps/replica:0/task:1")
self.assertAllEqual(v.variables[0].read_value().numpy(), [0, 1])
self.assertAllEqual(v.variables[1].read_value().numpy(), [2, 3])
def test_loading_partitions_no_max_rows(self):
"""Tests loading partitioned var as single slice with no valid max_rows."""
self._test_loading_variable_with_max_rows(
np_value=np.reshape(list(range(0, 36)), (9, 4)),
partitioner=partitioned_variables.fixed_size_partitioner(3),
max_rows_in_memory=-1)
def testSinglePartitionedVariable(self):
"""Ensures partitioned variables fail cleanly with freeze graph."""
checkpoint_prefix = os.path.join(self.get_temp_dir(), "saved_checkpoint")
checkpoint_state_name = "checkpoint_state"
input_graph_name = "input_graph.pb"
output_graph_name = "output_graph.pb"
# Create a graph with partition variables. When weights are partitioned into
# a single partition, the weights variable is followed by a identity ->
# identity (an additional identity node).
partitioner = partitioned_variables.fixed_size_partitioner(1)
with ops.Graph().as_default():
with variable_scope.variable_scope("part", partitioner=partitioner):
batch_size, height, width, depth = 5, 128, 128, 3
input1 = array_ops.zeros(
(batch_size, height, width, depth), name="input1")
input2 = array_ops.zeros(
(batch_size, height, width, depth), name="input2")
num_nodes = depth
filter1 = variable_scope.get_variable("filter", [num_nodes, num_nodes])
filter2 = array_ops.reshape(filter1, [1, 1, num_nodes, num_nodes])
conv = nn.conv2d(
input=input1, filter=filter2, strides=[1, 1, 1, 1], padding="SAME")
node = math_ops.add(conv, input2, name="test/add")
node = nn.relu6(node, name="test/relu6")
# Save graph and checkpoints.
sess = session.Session()
sess.run(variables.global_variables_initializer())
saver = saver_lib.Saver()
checkpoint_path = saver.save(
sess,
checkpoint_prefix,
global_step=0,
latest_filename=checkpoint_state_name)
graph_io.write_graph(sess.graph, self.get_temp_dir(), input_graph_name)
# Ensure this graph has partition variables.
self.assertTrue([
tensor.name.split(":")[0]
for op in sess.graph.get_operations()
for tensor in op.values()
if re.search(r"/part_\d+/", tensor.name)
])
# Test freezing graph doesn't make it crash.
output_node_names = "save/restore_all"
output_graph_path = os.path.join(self.get_temp_dir(), output_graph_name)
return_value = freeze_graph.freeze_graph_with_def_protos(
input_graph_def=sess.graph_def,
input_saver_def=None,
input_checkpoint=checkpoint_path,
output_node_names=output_node_names,
restore_op_name="save/restore_all", # default value
filename_tensor_name="save/Const:0", # default value
output_graph=output_graph_path,
clear_devices=False,
initializer_nodes="")
self.assertTrue(return_value, -1)
def __call__(self, inputs, state, scope=None):
"""Run one step of LSTM.
Args:
inputs: input Tensor, 2D, batch x num_units.
state: if `state_is_tuple` is False, this must be a state Tensor,
`2-D, batch x state_size`. If `state_is_tuple` is True, this must be a
tuple of state Tensors, both `2-D`, with column sizes `c_state` and
`m_state`.
scope: VariableScope for the created subgraph; defaults to "lstm_cell".
Returns:
A tuple containing:
- A `2-D, [batch x output_dim]`, Tensor representing the output of the
LSTM after reading `inputs` when previous state was `state`.
Here output_dim is:
num_proj if num_proj was set,
num_units otherwise.
- Tensor(s) representing the new state of LSTM after reading `inputs` when
the previous state was `state`. Same type and shape(s) as `state`.
Raises:
ValueError: If input size cannot be inferred from inputs via
static shape inference.
"""
num_proj = self._num_units if self._num_proj is None else self._num_proj
if self._state_is_tuple:
(c_prev, m_prev) = state
else:
c_prev = array_ops.slice(state, [0, 0], [-1, self._num_units])
m_prev = array_ops.slice(state, [0, self._num_units], [-1, num_proj])
dtype = inputs.dtype
input_size = inputs.get_shape().with_rank(2)[1]
if input_size.value is None:
raise ValueError("Could not infer input size from inputs.get_shape()[-1]")
with _checked_scope(self, scope or "lstm_cell",
initializer=self._initializer,
reuse=self._reuse) as unit_scope:
if self._num_unit_shards is not None:
unit_scope.set_partitioner(
partitioned_variables.fixed_size_partitioner(
self._num_unit_shards))
# i = input_gate, j = new_input, f = forget_gate, o = output_gate
lstm_matrix = _linear([inputs, m_prev], 4 * self._num_units, bias=True)
i, j, f, o = array_ops.split(
value=lstm_matrix, num_or_size_splits=4, axis=1)
# Diagonal connections
if self._use_peepholes:
with vs.variable_scope(unit_scope) as projection_scope:
if self._num_unit_shards is not None:
projection_scope.set_partitioner(None)
w_f_diag = vs.get_variable(
"w_f_diag", shape=[self._num_units], dtype=dtype)
w_i_diag = vs.get_variable(
"w_i_diag", shape=[self._num_units], dtype=dtype)
w_o_diag = vs.get_variable(
"w_o_diag", shape=[self._num_units], dtype=dtype)
if self._use_peepholes:
c = (sigmoid(f + self._forget_bias + w_f_diag * c_prev) * c_prev +
sigmoid(i + w_i_diag * c_prev) * self._activation(j))
else:
c = (sigmoid(f + self._forget_bias) * c_prev + sigmoid(i) *
self._activation(j))
if self._cell_clip is not None:
# pylint: disable=invalid-unary-operand-type
c = clip_ops.clip_by_value(c, -self._cell_clip, self._cell_clip)
# pylint: enable=invalid-unary-operand-type
if self._use_peepholes:
m = sigmoid(o + w_o_diag * c) * self._activation(c)
else:
m = sigmoid(o) * self._activation(c)
if self._num_proj is not None:
with vs.variable_scope("projection") as proj_scope:
if self._num_proj_shards is not None:
proj_scope.set_partitioner(
partitioned_variables.fixed_size_partitioner(
self._num_proj_shards))
m = _linear(m, self._num_proj, bias=False)
if self._proj_clip is not None:
# pylint: disable=invalid-unary-operand-type
m = clip_ops.clip_by_value(m, -self._proj_clip, self._proj_clip)
# pylint: enable=invalid-unary-operand-type
new_state = (LSTMStateTuple(c, m) if self._state_is_tuple else
array_ops.concat([c, m], 1))
return m, new_state
