-
Notifications
You must be signed in to change notification settings - Fork 346
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[feature]: add embedding parallel (#438)
* add parquet input in packed format and non-packed format * add support for embedding parallel, shard embedding across workers, based on horovod all2all and allreduce * refactor predictor into separate files and add parquet predictors * add script for custom ops build * add adam sparse optimizer * add fast and memory efficient auc implementation
- Loading branch information
1 parent
cb806b2
commit fddb73a
Showing
64 changed files
with
6,603 additions
and
705 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,2 +1,3 @@ | ||
| include easy_rec/python/ops/1.12/*.so* | ||
| include easy_rec/python/ops/1.15/*.so* | ||
| include easy_rec/python/ops/2.12/*.so* |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,245 @@ | ||
| # Copyright 2015 The TensorFlow Authors. All Rights Reserved. | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| # you may not use this file except in compliance with the License. | ||
| # You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software | ||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
| # ============================================================================== | ||
| """Adam for TensorFlow.""" | ||
| from __future__ import absolute_import | ||
| from __future__ import division | ||
| from __future__ import print_function | ||
|
|
||
| from tensorflow.python.eager import context | ||
| from tensorflow.python.framework import ops | ||
| from tensorflow.python.ops import array_ops | ||
| from tensorflow.python.ops import control_flow_ops | ||
| from tensorflow.python.ops import math_ops | ||
| from tensorflow.python.ops import resource_variable_ops | ||
| from tensorflow.python.ops import state_ops | ||
| from tensorflow.python.training import optimizer | ||
| from tensorflow.python.training import training_ops | ||
|
|
||
|
|
||
| class AdamOptimizerS(optimizer.Optimizer): | ||
| """Optimizer that implements the Adam algorithm. | ||
| References: | ||
| Adam - A Method for Stochastic Optimization: | ||
| [Kingma et al., 2015](https://arxiv.org/abs/1412.6980) | ||
| ([pdf](https://arxiv.org/pdf/1412.6980.pdf)) | ||
| """ | ||
|
|
||
| def __init__(self, | ||
| learning_rate=0.001, | ||
| beta1=0.9, | ||
| beta2=0.999, | ||
| epsilon=1e-8, | ||
| use_locking=False, | ||
| name='Adam'): | ||
| r"""Construct a new Adam optimizer. | ||
| Initialization: | ||
| $$m_0 := 0 \text{(Initialize initial 1st moment vector)}$$ | ||
| $$v_0 := 0 \text{(Initialize initial 2nd moment vector)}$$ | ||
| $$t := 0 \text{(Initialize timestep)}$$ | ||
| The update rule for `variable` with gradient `g` uses an optimization | ||
| described at the end of section 2 of the paper: | ||
| $$t := t + 1$$ | ||
| $$\text{lr}_t := \mathrm{learning_rate} * | ||
| \sqrt{1 - \beta_2^t} / (1 - \beta_1^t)$$ | ||
| $$m_t := \beta_1 * m_{t-1} + (1 - \beta_1) * g$$ | ||
| $$v_t := \beta_2 * v_{t-1} + (1 - \beta_2) * g * g$$ | ||
| $$\text{variable} := \text{variable} - | ||
| \text{lr}_t * m_t / (\sqrt{v_t} + \epsilon)$$ | ||
| The default value of 1e-8 for epsilon might not be a good default in | ||
| general. For example, when training an Inception network on ImageNet a | ||
| current good choice is 1.0 or 0.1. Note that since AdamOptimizerS uses the | ||
| formulation just before Section 2.1 of the Kingma and Ba paper rather than | ||
| the formulation in Algorithm 1, the "epsilon" referred to here is "epsilon | ||
| hat" in the paper. | ||
| The sparse implementation of this algorithm (used when the gradient is an | ||
| IndexedSlices object, typically because of `tf.gather` or an embedding | ||
| lookup in the forward pass) does apply momentum to variable slices even if | ||
| they were not used in the forward pass (meaning they have a gradient equal | ||
| to zero). Momentum decay (beta1) is also applied to the entire momentum | ||
| accumulator. This means that the sparse behavior is equivalent to the dense | ||
| behavior (in contrast to some momentum implementations which ignore momentum | ||
| unless a variable slice was actually used). | ||
| Args: | ||
| learning_rate: A Tensor or a floating point value. The learning rate. | ||
| beta1: A float value or a constant float tensor. The exponential decay | ||
| rate for the 1st moment estimates. | ||
| beta2: A float value or a constant float tensor. The exponential decay | ||
| rate for the 2nd moment estimates. | ||
| epsilon: A small constant for numerical stability. This epsilon is | ||
| "epsilon hat" in the Kingma and Ba paper (in the formula just before | ||
| Section 2.1), not the epsilon in Algorithm 1 of the paper. | ||
| use_locking: If True use locks for update operations. | ||
| name: Optional name for the operations created when applying gradients. | ||
| Defaults to "Adam". | ||
| @compatibility(eager) | ||
| When eager execution is enabled, `learning_rate`, `beta1`, `beta2`, and | ||
| `epsilon` can each be a callable that takes no arguments and returns the | ||
| actual value to use. This can be useful for changing these values across | ||
| different invocations of optimizer functions. | ||
| @end_compatibility | ||
| """ | ||
| super(AdamOptimizerS, self).__init__(use_locking, name) | ||
| self._lr = learning_rate | ||
| self._beta1 = beta1 | ||
| self._beta2 = beta2 | ||
| self._epsilon = epsilon | ||
|
|
||
| # Tensor versions of the constructor arguments, created in _prepare(). | ||
| self._lr_t = None | ||
| self._beta1_t = None | ||
| self._beta2_t = None | ||
| self._epsilon_t = None | ||
|
|
||
| def _get_beta_accumulators(self): | ||
| with ops.init_scope(): | ||
| if context.executing_eagerly(): | ||
| graph = None | ||
| else: | ||
| graph = ops.get_default_graph() | ||
| return (self._get_non_slot_variable('beta1_power', graph=graph), | ||
| self._get_non_slot_variable('beta2_power', graph=graph)) | ||
|
|
||
| def _create_slots(self, var_list): | ||
| # Create the beta1 and beta2 accumulators on the same device as the first | ||
| # variable. Sort the var_list to make sure this device is consistent across | ||
| # workers (these need to go on the same PS, otherwise some updates are | ||
| # silently ignored). | ||
| first_var = min(var_list, key=lambda x: x.name) | ||
| self._create_non_slot_variable( | ||
| initial_value=self._beta1, name='beta1_power', colocate_with=first_var) | ||
| self._create_non_slot_variable( | ||
| initial_value=self._beta2, name='beta2_power', colocate_with=first_var) | ||
|
|
||
| # Create slots for the first and second moments. | ||
| for v in var_list: | ||
| self._zeros_slot(v, 'm', self._name) | ||
| self._zeros_slot(v, 'v', self._name) | ||
|
|
||
| def _prepare(self): | ||
| lr = self._call_if_callable(self._lr) | ||
| beta1 = self._call_if_callable(self._beta1) | ||
| beta2 = self._call_if_callable(self._beta2) | ||
| epsilon = self._call_if_callable(self._epsilon) | ||
|
|
||
| self._lr_t = ops.convert_to_tensor(lr, name='learning_rate') | ||
| self._beta1_t = ops.convert_to_tensor(beta1, name='beta1') | ||
| self._beta2_t = ops.convert_to_tensor(beta2, name='beta2') | ||
| self._epsilon_t = ops.convert_to_tensor(epsilon, name='epsilon') | ||
|
|
||
| def _apply_dense(self, grad, var): | ||
| m = self.get_slot(var, 'm') | ||
| v = self.get_slot(var, 'v') | ||
| beta1_power, beta2_power = self._get_beta_accumulators() | ||
| return training_ops.apply_adam( | ||
| var, | ||
| m, | ||
| v, | ||
| math_ops.cast(beta1_power, var.dtype.base_dtype), | ||
| math_ops.cast(beta2_power, var.dtype.base_dtype), | ||
| math_ops.cast(self._lr_t, var.dtype.base_dtype), | ||
| math_ops.cast(self._beta1_t, var.dtype.base_dtype), | ||
| math_ops.cast(self._beta2_t, var.dtype.base_dtype), | ||
| math_ops.cast(self._epsilon_t, var.dtype.base_dtype), | ||
| grad, | ||
| use_locking=self._use_locking).op | ||
|
|
||
| def _resource_apply_dense(self, grad, var): | ||
| m = self.get_slot(var, 'm') | ||
| v = self.get_slot(var, 'v') | ||
| beta1_power, beta2_power = self._get_beta_accumulators() | ||
| return training_ops.resource_apply_adam( | ||
| var.handle, | ||
| m.handle, | ||
| v.handle, | ||
| math_ops.cast(beta1_power, grad.dtype.base_dtype), | ||
| math_ops.cast(beta2_power, grad.dtype.base_dtype), | ||
| math_ops.cast(self._lr_t, grad.dtype.base_dtype), | ||
| math_ops.cast(self._beta1_t, grad.dtype.base_dtype), | ||
| math_ops.cast(self._beta2_t, grad.dtype.base_dtype), | ||
| math_ops.cast(self._epsilon_t, grad.dtype.base_dtype), | ||
| grad, | ||
| use_locking=self._use_locking) | ||
|
|
||
| def _apply_sparse_shared(self, grad, var, indices, scatter_add): | ||
| beta1_power, beta2_power = self._get_beta_accumulators() | ||
| beta1_power = math_ops.cast(beta1_power, var.dtype.base_dtype) | ||
| beta2_power = math_ops.cast(beta2_power, var.dtype.base_dtype) | ||
| lr_t = math_ops.cast(self._lr_t, var.dtype.base_dtype) | ||
| beta1_t = math_ops.cast(self._beta1_t, var.dtype.base_dtype) | ||
| beta2_t = math_ops.cast(self._beta2_t, var.dtype.base_dtype) | ||
| epsilon_t = math_ops.cast(self._epsilon_t, var.dtype.base_dtype) | ||
| lr = (lr_t * math_ops.sqrt(1 - beta2_power) / (1 - beta1_power)) | ||
| # m_t = beta1 * m + (1 - beta1) * g_t | ||
| m = self.get_slot(var, 'm') | ||
| m_scaled_g_values = grad * (1 - beta1_t) | ||
| # m_t = state_ops.assign(m, m * beta1_t, use_locking=self._use_locking) | ||
| m_decay = array_ops.gather(m, indices) * beta1_t | ||
| m_part_n = m_scaled_g_values + m_decay | ||
| m_t = state_ops.scatter_update(m, indices, m_part_n) | ||
| # v_t = beta2 * v + (1 - beta2) * (g_t * g_t) | ||
| v = self.get_slot(var, 'v') | ||
| v_scaled_g_values = (grad * grad) * (1 - beta2_t) | ||
| v_decay = array_ops.gather(v, indices) * beta2_t | ||
| v_part_n = v_scaled_g_values + v_decay | ||
| v_t = state_ops.scatter_update(v, indices, v_part_n) | ||
| # v_sqrt = math_ops.sqrt(v_t) | ||
| # var_update = state_ops.assign_sub( | ||
| # var, lr * m_t / (v_sqrt + epsilon_t), use_locking=self._use_locking) | ||
| v_part_sqrt = math_ops.sqrt(v_part_n) | ||
| var_update = scatter_add(var, indices, | ||
| -lr * m_part_n / (v_part_sqrt + epsilon_t)) | ||
| return control_flow_ops.group(*[var_update, m_t, v_t]) | ||
|
|
||
| def _apply_sparse(self, grad, var): | ||
| return self._apply_sparse_shared( | ||
| grad.values, | ||
| var, | ||
| grad.indices, | ||
| lambda x, i, v: state_ops.scatter_add( # pylint: disable=g-long-lambda | ||
| x, | ||
| i, | ||
| v, | ||
| use_locking=self._use_locking)) | ||
|
|
||
| def _resource_scatter_add(self, x, i, v): | ||
| with ops.control_dependencies( | ||
| [resource_variable_ops.resource_scatter_add(x.handle, i, v)]): | ||
| return x.value() | ||
|
|
||
| def _resource_apply_sparse(self, grad, var, indices): | ||
| return self._apply_sparse_shared(grad, var, indices, | ||
| self._resource_scatter_add) | ||
|
|
||
| def _finish(self, update_ops, name_scope): | ||
| # Update the power accumulators. | ||
| with ops.control_dependencies(update_ops): | ||
| beta1_power, beta2_power = self._get_beta_accumulators() | ||
| with ops.colocate_with(beta1_power): | ||
| update_beta1 = beta1_power.assign( | ||
| beta1_power * self._beta1_t, use_locking=self._use_locking) | ||
| update_beta2 = beta2_power.assign( | ||
| beta2_power * self._beta2_t, use_locking=self._use_locking) | ||
| return control_flow_ops.group( | ||
| *update_ops + [update_beta1, update_beta2], name=name_scope) |
Oops, something went wrong.