def generate_send_gradients_op(self, feature_to_gradient_dict): """Send gradient to TPU embedding. Args: feature_to_gradient_dict: dict mapping feature names to gradient wrt activations. Returns: SendTPUEmbeddingGradients Op. Raises: RuntimeError: If `mode` is not `TRAINING`. """ if self._mode != TRAINING: raise RuntimeError('Only in training mode gradients need to ' 'be sent to TPU embedding; got mode {}.'.format( self._mode)) gradients = [] for table in self._table_to_features_dict: features = self._table_to_features_dict[table] table_gradients = [ feature_to_gradient_dict[feature] for feature in features ] concat_table_grads = array_ops.concat(table_gradients, axis=0) gradients.append(concat_table_grads) return tpu_ops.send_tpu_embedding_gradients( inputs=gradients, config=self.config_proto.SerializeToString())
def generate_send_gradients_op(self, gradient_multipliers=None): """Retrieve gradients from collections and send them to TPU embedding. Args: gradient_multipliers: None, or dict mapping table names to gradient multiplier Tensors. Returns: SendTPUEmbeddingGradients Op. Raises: ValueError: If required gradients have not been defined. RuntimeError: If `mode` is not `TRAINING`. """ if self._mode != TRAINING: raise RuntimeError('Only in training mode gradients need to ' 'be sent to TPU embedding; got mode {}.'.format( self._mode)) g = ops.get_default_graph() gradients = list() for table_id, table in enumerate(self._table_to_config_dict): table_gradients = g.get_collection( 'tpu_embedding_gradients_table_%d' % table_id) if any(gradient is None for gradient in table_gradients): raise ValueError( 'Table {}/{} has undefined gradients: this is probably because the ' 'model asked TPUEmbedding to compute activations that were not ' 'used.'.format(table_id, table)) concat_table_grads = array_ops.concat(table_gradients, axis=0) if gradient_multipliers is not None: concat_table_grads *= gradient_multipliers[table.name] gradients.append(concat_table_grads) return tpu_ops.send_tpu_embedding_gradients( inputs=gradients, config=self.config_proto.SerializeToString())
def generate_send_gradients_op(self, gradient_multipliers=None): """Retrieve gradients from collections and send them to TPU embedding. Args: gradient_multipliers: None, or dict mapping table names to gradient multiplier Tensors. Returns: SendTPUEmbeddingGradients Op. Raises: ValueError: If required gradients have not been defined. RuntimeError: If `mode` is not `TRAINING`. """ if self._mode != TRAINING: raise RuntimeError('Only in training mode gradients need to ' 'be sent to TPU embedding; got mode {}.' .format(self._mode)) g = ops.get_default_graph() gradients = list() for table_id, table in enumerate(self._table_to_config_dict): table_gradients = g.get_collection( 'tpu_embedding_gradients_table_%d' % table_id) if any(gradient is None for gradient in table_gradients): raise ValueError( 'Table {}/{} has undefined gradients: this is probably because the ' 'model asked TPUEmbedding to compute activations that were not ' 'used.'.format(table_id, table)) concat_table_grads = array_ops.concat(table_gradients, axis=0) if gradient_multipliers is not None: concat_table_grads *= gradient_multipliers[table.name] gradients.append(concat_table_grads) return tpu_ops.send_tpu_embedding_gradients( inputs=gradients, config=self.config_proto.SerializeToString())