def testTransform(self): # This tests all combinations of: # - ids rank 0, 1, >1 # - params sharded/unsharded # It always applies max_norm. np.random.seed(8) l2_norm = 2. with self.test_session(): # Param values are in [l2_norm, l2_norm+1) so it will always clip. params = np.random.rand(6, 3) + l2_norm params_norm = l2_norm * params / np.sqrt( np.sum(params * params, axis=1, keepdims=True)) # Compute the norm of each embedding. This will change the embedding # rank to 0. params_norm = np.linalg.norm(params_norm, axis=1) transform = lambda x: linalg_ops.norm(x, axis=1) for ids_shape in (), (3), (4, 3), (2, 3, 4): # Test ids rank 0, 1, 2, 3. ids = np.random.randint( params.shape[0], size=np.prod(ids_shape, dtype=np.int64)).reshape(ids_shape) # Compare nonsharded to gather. simple = embedding_ops._embedding_lookup_and_transform( params, ids, max_norm=l2_norm, transform_fn=transform).eval() self.assertAllClose(simple, array_ops.gather(params_norm, ids).eval()) # Run a few different sharded versions. for procs in 1, 2, 3: stride = procs * math_ops.range(params.shape[0] // procs) split_params = [ array_ops.gather(params, stride + p) for p in xrange(procs) ] sharded = embedding_ops._embedding_lookup_and_transform( split_params, ids, max_norm=l2_norm, transform_fn=transform).eval() self.assertAllEqual(simple, sharded)
def testTransform(self): # This tests all combinations of: # - ids rank 0, 1, >1 # - params sharded/unsharded # It always applies max_norm. np.random.seed(8) l2_norm = 2. with self.cached_session(): # Param values are in [l2_norm, l2_norm+1) so it will always clip. params = np.random.rand(6, 3) + l2_norm params_norm = l2_norm * params / np.sqrt( np.sum(params * params, axis=1, keepdims=True)) # Compute the norm of each embedding. This will change the embedding # rank to 0. params_norm = np.linalg.norm(params_norm, axis=1) transform = lambda x: linalg_ops.norm(x, axis=1) for ids_shape in (), (3), (4, 3), (2, 3, 4): # Test ids rank 0, 1, 2, 3. ids = np.random.randint( params.shape[0], size=np.prod(ids_shape, dtype=np.int64)).reshape(ids_shape) # Compare nonsharded to gather. simple = embedding_ops._embedding_lookup_and_transform( params, ids, max_norm=l2_norm, transform_fn=transform).eval() self.assertAllClose(simple, array_ops.gather(params_norm, ids).eval()) # Run a few different sharded versions. for procs in 1, 2, 3: stride = procs * math_ops.range(params.shape[0] // procs) split_params = [ array_ops.gather(params, stride + p) for p in xrange(procs) ] sharded = embedding_ops._embedding_lookup_and_transform( split_params, ids, max_norm=l2_norm, transform_fn=transform).eval() self.assertAllEqual(simple, sharded)
def testTransform(self): # This tests all combinations of: # - ids rank 0, 1, >1 # - params sharded/unsharded # It always applies max_norm. np.random.seed(8) l2_norm = 2. with self.cached_session(): # Param values are in [l2_norm, l2_norm+1) so it will always clip. params = np.random.rand(6, 3) + l2_norm params_norm = l2_norm * params / np.sqrt( np.sum(params * params, axis=1, keepdims=True)) # Compute the norm of each embedding. This will change the embedding # rank to 0. params_norm = np.linalg.norm(params_norm, axis=1) transform = lambda x: linalg_ops.norm(x, axis=1) for ids_shape in (), (3), (4, 3), (2, 3, 4): # Test ids rank 0, 1, 2, 3. ids = np.random.randint( params.shape[0], size=np.prod(ids_shape, dtype=np.int64)).reshape(ids_shape) # Compare nonsharded to gather. simple = embedding_ops._embedding_lookup_and_transform( params, ids, max_norm=l2_norm, transform_fn=transform).eval() self.assertAllClose(simple, array_ops.gather(params_norm, ids)) # Run a few different sharded versions. for procs in 1, 2, 3: stride = procs * math_ops.range(params.shape[0] // procs) split_params = [ array_ops.gather(params, stride + p) for p in xrange(procs) ] sharded = embedding_ops._embedding_lookup_and_transform( split_params, ids, max_norm=l2_norm, transform_fn=transform).eval() # assertAllClose is used here as different implementations of sqrt may # be used to compute each of the values being compared. For example, # on AVX512 builds the embedding operation makes use of Eigen's fast # vectorized square root algorithm for doubles. These different # implementations of sqrt are not guaranteed to produce exactly the # same results. Therefore, an exact comparison cannot be made. self.assertAllClose(simple, sharded)
def testTransform(self): # This tests all combinations of: # - ids rank 0, 1, >1 # - params sharded/unsharded # It always applies max_norm. np.random.seed(8) l2_norm = 2. with self.cached_session(): # Param values are in [l2_norm, l2_norm+1) so it will always clip. params = np.random.rand(6, 3) + l2_norm params_norm = l2_norm * params / np.sqrt( np.sum(params * params, axis=1, keepdims=True)) # Compute the norm of each embedding. This will change the embedding # rank to 0. params_norm = np.linalg.norm(params_norm, axis=1) transform = lambda x: linalg_ops.norm(x, axis=1) for ids_shape in (), (3), (4, 3), (2, 3, 4): # Test ids rank 0, 1, 2, 3. ids = np.random.randint( params.shape[0], size=np.prod(ids_shape, dtype=np.int64)).reshape(ids_shape) # Compare nonsharded to gather. simple = embedding_ops._embedding_lookup_and_transform( params, ids, max_norm=l2_norm, transform_fn=transform).eval() self.assertAllClose(simple, array_ops.gather(params_norm, ids).eval()) # Run a few different sharded versions. for procs in 1, 2, 3: stride = procs * math_ops.range(params.shape[0] // procs) split_params = [ array_ops.gather(params, stride + p) for p in xrange(procs) ] sharded = embedding_ops._embedding_lookup_and_transform( split_params, ids, max_norm=l2_norm, transform_fn=transform).eval() # assertAllClose is used here as different implementations of sqrt may # be used to compute each of the values being compared. For example, # on AVX512 builds the embedding operation makes use of Eigen's fast # vectorized square root algorithm for doubles. These different # implementations of sqrt are not guaranteed to produce exactly the # same results. Therefore, an exact comparison cannot be made. self.assertAllClose(simple, sharded)
def _rank_resample(weights, biases, inputs, sampled_values, num_resampled, resampling_temperature, partition_strategy): """A helper function for rank_sampled_softmax_loss. This computes, for each i in `sampled_values`, log(sum_j exp((w_i * x_j + b_i) / resampling_temperature)) where w_i, b_i are the weight and bias of the i-th class, respectively, and j ranges over the rows of `inputs`. For efficiency, we rearrange the computation to log(sum_j exp(w_i * (x_j / resampling_temperature))) + b_i / resampling_temperature. This translates to the following batched computation using tensorflow ops: reduce_logsumexp(matmul(embeddings, transpose(inputs / resampling_temperature))) + biases / resampling_temperature The computation of the first term is colocated with the embeddings using `transform_fn` in `embedding_ops._embedding_lookup_and_transform`. The second term, not the bottleneck, is computed at the worker. Args: weights: From `rank_sampled_softmax_loss`. biases: From `rank_sampled_softmax_loss`. inputs: From `rank_sampled_softmax_loss`. sampled_values: A tuple of (`sampled_candidates`, `true_expected_count`, `sampled_expected_count`) returned by a `*_candidate_sampler` function. num_resampled: An `int`. This many values are selected from `sampled_values` using the adaptive resampling algorithm. The caller must ensure that `num_resampled` is less than the size of `sampled_values`. resampling_temperature: A scalar `Tensor` with the temperature parameter for the adaptive resampling algorithm. partition_strategy: From `rank_sampled_softmax_loss`. Returns: A tuple of (`resampled_candidates`, `true_expected_count`, `resampled_expected_count`), similar to `sampled_values` but sampled down to `num_resampled` values. """ # This code supports passing a Tensor for num_resampled, but since it is only # called with an int, that's what we specify in the arg list. If this # function is ever externalized, we should change the doc to support Tensor. sampled, true_expected_count, sampled_expected_count = sampled_values sampled = math_ops.cast(array_ops.stop_gradient(sampled), dtypes.int64) true_expected_count = array_ops.stop_gradient(true_expected_count) sampled_expected_count = array_ops.stop_gradient(sampled_expected_count) reweighted_inputs = inputs / resampling_temperature def logsumexp_logit(embeddings): return math_ops.reduce_logsumexp( math_ops.matmul(embeddings, reweighted_inputs, transpose_b=True), axis=1, keepdims=False) # Calling this protected form of embedding_lookup allows co-locating # the logsumexp computation with the partitioned weights, which yields # a large speedup in practice. sampled_logits = embedding_ops._embedding_lookup_and_transform( # pylint: disable=protected-access weights, sampled, partition_strategy, transform_fn=logsumexp_logit) sampled_b = array_ops.reshape( embedding_ops.embedding_lookup(biases, sampled, partition_strategy), [-1]) sampled_logits += sampled_b / resampling_temperature _, resampled_indices = nn.top_k(sampled_logits, k=num_resampled, sorted=False) resampled = array_ops.gather(sampled, indices=resampled_indices) resampled_expected_count = array_ops.gather( sampled_expected_count, indices=resampled_indices) return resampled, true_expected_count, resampled_expected_count
def _rank_resample(weights, biases, inputs, sampled_values, num_resampled, resampling_temperature, partition_strategy): """A helper function for rank_sampled_softmax_loss. This computes, for each i in `sampled_values`, log(sum_j exp((w_i * x_j + b_i) / resampling_temperature)) where w_i, b_i are the weight and bias of the i-th class, respectively, and j ranges over the rows of `inputs`. For efficiency, we rearrange the computation to log(sum_j exp(w_i * (x_j / resampling_temperature))) + b_i / resampling_temperature. This translates to the following batched computation using tensorflow ops: reduce_logsumexp(matmul(embeddings, transpose(inputs / resampling_temperature))) + biases / resampling_temperature The computation of the first term is colocated with the embeddings using `transform_fn` in `embedding_ops._embedding_lookup_and_transform`. The second term, not the bottleneck, is computed at the worker. Args: weights: From `rank_sampled_softmax_loss`. biases: From `rank_sampled_softmax_loss`. inputs: From `rank_sampled_softmax_loss`. sampled_values: A tuple of (`sampled_candidates`, `true_expected_count`, `sampled_expected_count`) returned by a `*_candidate_sampler` function. num_resampled: An `int`. This many values are selected from `sampled_values` using the adaptive resampling algorithm. The caller must ensure that `num_resampled` is less than the size of `sampled_values`. resampling_temperature: A scalar `Tensor` with the temperature parameter for the adaptive resampling algorithm. partition_strategy: From `rank_sampled_softmax_loss`. Returns: A tuple of (`resampled_candidates`, `true_expected_count`, `resampled_expected_count`), similar to `sampled_values` but sampled down to `num_resampled` values. """ # This code supports passing a Tensor for num_resampled, but since it is only # called with an int, that's what we specify in the arg list. If this # function is ever externalized, we should change the doc to support Tensor. sampled, true_expected_count, sampled_expected_count = sampled_values sampled = math_ops.cast(array_ops.stop_gradient(sampled), dtypes.int64) true_expected_count = array_ops.stop_gradient(true_expected_count) sampled_expected_count = array_ops.stop_gradient(sampled_expected_count) reweighted_inputs = inputs / resampling_temperature def logsumexp_logit(embeddings): return math_ops.reduce_logsumexp(math_ops.matmul(embeddings, reweighted_inputs, transpose_b=True), axis=1, keepdims=False) # Calling this protected form of embedding_lookup allows co-locating # the logsumexp computation with the partitioned weights, which yields # a large speedup in practice. sampled_logits = embedding_ops._embedding_lookup_and_transform( # pylint: disable=protected-access weights, sampled, partition_strategy, transform_fn=logsumexp_logit) sampled_b = array_ops.reshape( embedding_ops.embedding_lookup(biases, sampled, partition_strategy), [-1]) sampled_logits += sampled_b / resampling_temperature _, resampled_indices = nn.top_k(sampled_logits, k=num_resampled, sorted=False) resampled = array_ops.gather(sampled, indices=resampled_indices) resampled_expected_count = array_ops.gather(sampled_expected_count, indices=resampled_indices) return resampled, true_expected_count, resampled_expected_count