775 lines
35 KiB
Python
775 lines
35 KiB
Python
|
# Copyright 2016 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.
|
||
|
|
# ==============================================================================
|
||
|
|
"""Clustering Operations."""
|
||
|
|
|
||
|
|
from tensorflow.python.framework import constant_op
|
||
|
|
from tensorflow.python.framework import dtypes
|
||
|
|
from tensorflow.python.framework import ops
|
||
|
|
from tensorflow.python.framework import random_seed as random_seed_ops
|
||
|
|
from tensorflow.python.ops import array_ops
|
||
|
|
from tensorflow.python.ops import check_ops
|
||
|
|
from tensorflow.python.ops import cond
|
||
|
|
from tensorflow.python.ops import control_flow_ops
|
||
|
|
from tensorflow.python.ops import gen_clustering_ops
|
||
|
|
from tensorflow.python.ops import math_ops
|
||
|
|
from tensorflow.python.ops import nn_impl
|
||
|
|
from tensorflow.python.ops import random_ops
|
||
|
|
from tensorflow.python.ops import state_ops
|
||
|
|
from tensorflow.python.ops import variable_v1
|
||
|
|
from tensorflow.python.ops import while_loop
|
||
|
|
from tensorflow.python.ops.embedding_ops import embedding_lookup
|
||
|
|
# go/tf-wildcard-import
|
||
|
|
# pylint: disable=wildcard-import
|
||
|
|
from tensorflow.python.ops.gen_clustering_ops import *
|
||
|
|
# pylint: enable=wildcard-import
|
||
|
|
|
||
|
|
# Euclidean distance between vectors U and V is defined as \\(||U - V||_F\\)
|
||
|
|
# which is the square root of the sum of the absolute squares of the elements
|
||
|
|
# difference.
|
||
|
|
SQUARED_EUCLIDEAN_DISTANCE = 'squared_euclidean'
|
||
|
|
# Cosine distance between vectors U and V is defined as
|
||
|
|
# \\(1 - (U \dot V) / (||U||_F ||V||_F)\\)
|
||
|
|
COSINE_DISTANCE = 'cosine'
|
||
|
|
|
||
|
|
RANDOM_INIT = 'random'
|
||
|
|
KMEANS_PLUS_PLUS_INIT = 'kmeans_plus_plus'
|
||
|
|
KMC2_INIT = 'kmc2'
|
||
|
|
|
||
|
|
CLUSTERS_VAR_NAME = 'clusters'
|
||
|
|
|
||
|
|
|
||
|
|
class KMeans:
|
||
|
|
"""Creates the graph for k-means clustering."""
|
||
|
|
|
||
|
|
def __init__(self,
|
||
|
|
inputs,
|
||
|
|
num_clusters,
|
||
|
|
initial_clusters=RANDOM_INIT,
|
||
|
|
distance_metric=SQUARED_EUCLIDEAN_DISTANCE,
|
||
|
|
use_mini_batch=False,
|
||
|
|
mini_batch_steps_per_iteration=1,
|
||
|
|
random_seed=0,
|
||
|
|
kmeans_plus_plus_num_retries=2,
|
||
|
|
kmc2_chain_length=200):
|
||
|
|
"""Creates an object for generating KMeans clustering graph.
|
||
|
|
|
||
|
|
This class implements the following variants of K-means algorithm:
|
||
|
|
|
||
|
|
If use_mini_batch is False, it runs standard full batch K-means. Each step
|
||
|
|
runs a single iteration of K-Means. This step can be run sharded across
|
||
|
|
multiple workers by passing a list of sharded inputs to this class. Note
|
||
|
|
however that a single step needs to process the full input at once.
|
||
|
|
|
||
|
|
If use_mini_batch is True, it runs a generalization of the mini-batch
|
||
|
|
K-means algorithm. It runs multiple iterations, where each iteration is
|
||
|
|
composed of mini_batch_steps_per_iteration steps. Two copies of cluster
|
||
|
|
centers are maintained: one that is updated at the end of each iteration,
|
||
|
|
and one that is updated every step. The first copy is used to compute
|
||
|
|
cluster allocations for each step, and for inference, while the second copy
|
||
|
|
is the one updated each step using the mini-batch update rule. After each
|
||
|
|
iteration is complete, this second copy is copied back the first copy.
|
||
|
|
|
||
|
|
Note that for use_mini_batch=True, when mini_batch_steps_per_iteration=1,
|
||
|
|
the algorithm reduces to the standard mini-batch algorithm. Also by setting
|
||
|
|
mini_batch_steps_per_iteration = num_inputs / batch_size, the algorithm
|
||
|
|
becomes an asynchronous version of the full-batch algorithm. Note however
|
||
|
|
that there is no guarantee by this implementation that each input is seen
|
||
|
|
exactly once per iteration. Also, different updates are applied
|
||
|
|
asynchronously without locking. So this asynchronous version may not behave
|
||
|
|
exactly like a full-batch version.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
inputs: An input tensor or list of input tensors. It is assumed that the
|
||
|
|
data points have been previously randomly permuted.
|
||
|
|
num_clusters: An integer tensor specifying the number of clusters. This
|
||
|
|
argument is ignored if initial_clusters is a tensor or numpy array.
|
||
|
|
initial_clusters: Specifies the clusters used during initialization. One
|
||
|
|
of the following: - a tensor or numpy array with the initial cluster
|
||
|
|
centers. - a function f(inputs, k) that returns up to k centers from
|
||
|
|
`inputs`.
|
||
|
|
- "random": Choose centers randomly from `inputs`.
|
||
|
|
- "kmeans_plus_plus": Use kmeans++ to choose centers from `inputs`.
|
||
|
|
- "kmc2": Use the fast k-MC2 algorithm to choose centers from `inputs`.
|
||
|
|
In the last three cases, one batch of `inputs` may not yield
|
||
|
|
`num_clusters` centers, in which case initialization will require
|
||
|
|
multiple batches until enough centers are chosen. In the case of
|
||
|
|
"random" or "kmeans_plus_plus", if the input size is <= `num_clusters`
|
||
|
|
then the entire batch is chosen to be cluster centers.
|
||
|
|
distance_metric: Distance metric used for clustering. Supported options:
|
||
|
|
"squared_euclidean", "cosine".
|
||
|
|
use_mini_batch: If true, use the mini-batch k-means algorithm. Else assume
|
||
|
|
full batch.
|
||
|
|
mini_batch_steps_per_iteration: Number of steps after which the updated
|
||
|
|
cluster centers are synced back to a master copy.
|
||
|
|
random_seed: Seed for PRNG used to initialize seeds.
|
||
|
|
kmeans_plus_plus_num_retries: For each point that is sampled during
|
||
|
|
kmeans++ initialization, this parameter specifies the number of
|
||
|
|
additional points to draw from the current distribution before selecting
|
||
|
|
the best. If a negative value is specified, a heuristic is used to
|
||
|
|
sample O(log(num_to_sample)) additional points.
|
||
|
|
kmc2_chain_length: Determines how many candidate points are used by the
|
||
|
|
k-MC2 algorithm to produce one new cluster centers. If a (mini-)batch
|
||
|
|
contains less points, one new cluster center is generated from the
|
||
|
|
(mini-)batch.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
ValueError: An invalid argument was passed to initial_clusters or
|
||
|
|
distance_metric.
|
||
|
|
"""
|
||
|
|
initialization_algorithms = [RANDOM_INIT, KMEANS_PLUS_PLUS_INIT, KMC2_INIT]
|
||
|
|
if isinstance(initial_clusters,
|
||
|
|
str) and initial_clusters not in initialization_algorithms:
|
||
|
|
raise ValueError(
|
||
|
|
f'Unsupported initialization algorithm `{initial_clusters}`,'
|
||
|
|
f'must be one of `{initialization_algorithms}`.')
|
||
|
|
|
||
|
|
distance_metrics = [SQUARED_EUCLIDEAN_DISTANCE, COSINE_DISTANCE]
|
||
|
|
if distance_metric not in distance_metrics:
|
||
|
|
raise ValueError(f'Unsupported distance metric `{distance_metric}`,'
|
||
|
|
f'must be one of `{distance_metrics}`.')
|
||
|
|
self._inputs = inputs if isinstance(inputs, list) else [inputs]
|
||
|
|
self._num_clusters = num_clusters
|
||
|
|
self._initial_clusters = initial_clusters
|
||
|
|
self._distance_metric = distance_metric
|
||
|
|
self._use_mini_batch = use_mini_batch
|
||
|
|
self._mini_batch_steps_per_iteration = int(mini_batch_steps_per_iteration)
|
||
|
|
self._seed = random_seed_ops.get_seed(random_seed)[0]
|
||
|
|
self._kmeans_plus_plus_num_retries = kmeans_plus_plus_num_retries
|
||
|
|
self._kmc2_chain_length = kmc2_chain_length
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def _distance_graph(cls, inputs, clusters, distance_metric):
|
||
|
|
"""Computes distance between each input and each cluster center.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
inputs: list of input Tensors.
|
||
|
|
clusters: cluster Tensor.
|
||
|
|
distance_metric: distance metric used for clustering
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
list of Tensors, where each element corresponds to each element in inputs.
|
||
|
|
The value is the distance of each row to all the cluster centers.
|
||
|
|
Currently only Euclidean distance and cosine distance are supported.
|
||
|
|
"""
|
||
|
|
assert isinstance(inputs, list)
|
||
|
|
if distance_metric == SQUARED_EUCLIDEAN_DISTANCE:
|
||
|
|
return cls._compute_euclidean_distance(inputs, clusters)
|
||
|
|
elif distance_metric == COSINE_DISTANCE:
|
||
|
|
return cls._compute_cosine_distance(
|
||
|
|
inputs, clusters, inputs_normalized=True)
|
||
|
|
else:
|
||
|
|
assert False, str(distance_metric)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def _compute_euclidean_distance(cls, inputs, clusters):
|
||
|
|
"""Computes Euclidean distance between each input and each cluster center.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
inputs: list of input Tensors.
|
||
|
|
clusters: cluster Tensor.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
list of Tensors, where each element corresponds to each element in inputs.
|
||
|
|
The value is the distance of each row to all the cluster centers.
|
||
|
|
"""
|
||
|
|
output = []
|
||
|
|
for inp in inputs:
|
||
|
|
with ops.colocate_with(inp, ignore_existing=True):
|
||
|
|
# Computes Euclidean distance. Note the first and third terms are
|
||
|
|
# broadcast additions.
|
||
|
|
squared_distance = (
|
||
|
|
math_ops.reduce_sum(math_ops.square(inp), 1, keepdims=True) -
|
||
|
|
2 * math_ops.matmul(inp, clusters, transpose_b=True) +
|
||
|
|
array_ops.transpose(
|
||
|
|
math_ops.reduce_sum(
|
||
|
|
math_ops.square(clusters), 1, keepdims=True)))
|
||
|
|
output.append(squared_distance)
|
||
|
|
|
||
|
|
return output
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def _compute_cosine_distance(cls, inputs, clusters, inputs_normalized=True):
|
||
|
|
"""Computes cosine distance between each input and each cluster center.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
inputs: list of input Tensor.
|
||
|
|
clusters: cluster Tensor
|
||
|
|
inputs_normalized: if True, it assumes that inp and clusters are
|
||
|
|
normalized and computes the dot product which is equivalent to the
|
||
|
|
cosine distance. Else it L2 normalizes the inputs first.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
list of Tensors, where each element corresponds to each element in inp.
|
||
|
|
The value is the distance of each row to all the cluster centers.
|
||
|
|
"""
|
||
|
|
output = []
|
||
|
|
if not inputs_normalized:
|
||
|
|
with ops.colocate_with(clusters, ignore_existing=True):
|
||
|
|
clusters = nn_impl.l2_normalize(clusters, axis=1)
|
||
|
|
for inp in inputs:
|
||
|
|
with ops.colocate_with(inp, ignore_existing=True):
|
||
|
|
if not inputs_normalized:
|
||
|
|
inp = nn_impl.l2_normalize(inp, axis=1)
|
||
|
|
output.append(1 - math_ops.matmul(inp, clusters, transpose_b=True))
|
||
|
|
return output
|
||
|
|
|
||
|
|
def _infer_graph(self, inputs, clusters):
|
||
|
|
"""Maps input to closest cluster and the score.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
inputs: list of input Tensors.
|
||
|
|
clusters: Tensor of cluster centers.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
List of tuple, where each value in tuple corresponds to a value in inp.
|
||
|
|
The tuple has following three elements:
|
||
|
|
all_scores: distance of each input to each cluster center.
|
||
|
|
score: distance of each input to closest cluster center.
|
||
|
|
cluster_idx: index of cluster center closest to the corresponding input.
|
||
|
|
"""
|
||
|
|
assert isinstance(inputs, list)
|
||
|
|
# Pairwise distances are used only by transform(). In all other cases, this
|
||
|
|
# sub-graph is not evaluated.
|
||
|
|
scores = self._distance_graph(inputs, clusters, self._distance_metric)
|
||
|
|
output = []
|
||
|
|
if (self._distance_metric == COSINE_DISTANCE and
|
||
|
|
not self._clusters_l2_normalized()):
|
||
|
|
# The cosine distance between normalized vectors x and y is the same as
|
||
|
|
# 2 * squared_euclidean_distance. We are using this fact and reusing the
|
||
|
|
# nearest_neighbors op.
|
||
|
|
# TODO(ands): Support COSINE distance in nearest_neighbors and remove
|
||
|
|
# this.
|
||
|
|
with ops.colocate_with(clusters, ignore_existing=True):
|
||
|
|
clusters = nn_impl.l2_normalize(clusters, axis=1)
|
||
|
|
for inp, score in zip(inputs, scores):
|
||
|
|
with ops.colocate_with(inp, ignore_existing=True):
|
||
|
|
(indices,
|
||
|
|
distances) = gen_clustering_ops.nearest_neighbors(inp, clusters, 1)
|
||
|
|
if self._distance_metric == COSINE_DISTANCE:
|
||
|
|
distances *= 0.5
|
||
|
|
output.append(
|
||
|
|
(score, array_ops.squeeze(distances,
|
||
|
|
[-1]), array_ops.squeeze(indices, [-1])))
|
||
|
|
return zip(*output)
|
||
|
|
|
||
|
|
def _clusters_l2_normalized(self):
|
||
|
|
"""Returns True if clusters centers are kept normalized."""
|
||
|
|
return (self._distance_metric == COSINE_DISTANCE and
|
||
|
|
(not self._use_mini_batch or
|
||
|
|
self._mini_batch_steps_per_iteration > 1))
|
||
|
|
|
||
|
|
def _create_variables(self, num_clusters):
|
||
|
|
"""Creates variables.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
num_clusters: an integer Tensor providing the number of clusters.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Tuple with following elements:
|
||
|
|
- cluster_centers: a Tensor for storing cluster centers
|
||
|
|
- cluster_centers_initialized: bool Variable indicating whether clusters
|
||
|
|
are initialized.
|
||
|
|
- cluster_counts: a Tensor for storing counts of points assigned to this
|
||
|
|
cluster. This is used by mini-batch training.
|
||
|
|
- cluster_centers_updated: Tensor representing copy of cluster centers
|
||
|
|
that are updated every step.
|
||
|
|
- update_in_steps: numbers of steps left before we sync
|
||
|
|
cluster_centers_updated back to cluster_centers.
|
||
|
|
"""
|
||
|
|
init_value = array_ops.placeholder_with_default([], shape=None)
|
||
|
|
cluster_centers = variable_v1.VariableV1(
|
||
|
|
init_value, name=CLUSTERS_VAR_NAME, validate_shape=False)
|
||
|
|
cluster_centers_initialized = variable_v1.VariableV1(
|
||
|
|
False, dtype=dtypes.bool, name='initialized')
|
||
|
|
|
||
|
|
if self._use_mini_batch and self._mini_batch_steps_per_iteration > 1:
|
||
|
|
# Copy of cluster centers actively updated each step according to
|
||
|
|
# mini-batch update rule.
|
||
|
|
cluster_centers_updated = variable_v1.VariableV1(
|
||
|
|
init_value, name='clusters_updated', validate_shape=False)
|
||
|
|
# How many steps till we copy the updated clusters to cluster_centers.
|
||
|
|
update_in_steps = variable_v1.VariableV1(
|
||
|
|
self._mini_batch_steps_per_iteration,
|
||
|
|
dtype=dtypes.int64,
|
||
|
|
name='update_in_steps')
|
||
|
|
# Count of points assigned to cluster_centers_updated.
|
||
|
|
cluster_counts = variable_v1.VariableV1(
|
||
|
|
array_ops.zeros([num_clusters], dtype=dtypes.int64))
|
||
|
|
else:
|
||
|
|
cluster_centers_updated = cluster_centers
|
||
|
|
update_in_steps = None
|
||
|
|
cluster_counts = (
|
||
|
|
variable_v1.VariableV1(
|
||
|
|
array_ops.ones([num_clusters], dtype=dtypes.int64))
|
||
|
|
if self._use_mini_batch else None)
|
||
|
|
return (cluster_centers, cluster_centers_initialized, cluster_counts,
|
||
|
|
cluster_centers_updated, update_in_steps)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def _l2_normalize_data(cls, inputs):
|
||
|
|
"""Normalized the input data."""
|
||
|
|
output = []
|
||
|
|
for inp in inputs:
|
||
|
|
with ops.colocate_with(inp, ignore_existing=True):
|
||
|
|
output.append(nn_impl.l2_normalize(inp, dim=1))
|
||
|
|
return output
|
||
|
|
|
||
|
|
def training_graph(self):
|
||
|
|
"""Generate a training graph for kmeans algorithm.
|
||
|
|
|
||
|
|
This returns, among other things, an op that chooses initial centers
|
||
|
|
(init_op), a boolean variable that is set to True when the initial centers
|
||
|
|
are chosen (cluster_centers_initialized), and an op to perform either an
|
||
|
|
entire Lloyd iteration or a mini-batch of a Lloyd iteration (training_op).
|
||
|
|
The caller should use these components as follows. A single worker should
|
||
|
|
execute init_op multiple times until cluster_centers_initialized becomes
|
||
|
|
True. Then multiple workers may execute training_op any number of times.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
A tuple consisting of:
|
||
|
|
all_scores: A matrix (or list of matrices) of dimensions (num_input,
|
||
|
|
num_clusters) where the value is the distance of an input vector and a
|
||
|
|
cluster center.
|
||
|
|
cluster_idx: A vector (or list of vectors). Each element in the vector
|
||
|
|
corresponds to an input row in 'inp' and specifies the cluster id
|
||
|
|
corresponding to the input.
|
||
|
|
scores: Similar to cluster_idx but specifies the distance to the
|
||
|
|
assigned cluster instead.
|
||
|
|
cluster_centers_initialized: scalar indicating whether clusters have been
|
||
|
|
initialized.
|
||
|
|
init_op: an op to initialize the clusters.
|
||
|
|
training_op: an op that runs an iteration of training.
|
||
|
|
"""
|
||
|
|
# Implementation of kmeans.
|
||
|
|
if (isinstance(self._initial_clusters, str) or
|
||
|
|
callable(self._initial_clusters)):
|
||
|
|
initial_clusters = self._initial_clusters
|
||
|
|
num_clusters = ops.convert_to_tensor(self._num_clusters)
|
||
|
|
else:
|
||
|
|
initial_clusters = ops.convert_to_tensor(self._initial_clusters)
|
||
|
|
num_clusters = array_ops.shape(initial_clusters)[0]
|
||
|
|
|
||
|
|
inputs = self._inputs
|
||
|
|
(cluster_centers_var, cluster_centers_initialized, total_counts,
|
||
|
|
cluster_centers_updated,
|
||
|
|
update_in_steps) = self._create_variables(num_clusters)
|
||
|
|
init_op = _InitializeClustersOpFactory(
|
||
|
|
self._inputs, num_clusters, initial_clusters, self._distance_metric,
|
||
|
|
self._seed, self._kmeans_plus_plus_num_retries, self._kmc2_chain_length,
|
||
|
|
cluster_centers_var, cluster_centers_updated,
|
||
|
|
cluster_centers_initialized).op()
|
||
|
|
cluster_centers = cluster_centers_var
|
||
|
|
|
||
|
|
if self._distance_metric == COSINE_DISTANCE:
|
||
|
|
inputs = self._l2_normalize_data(inputs)
|
||
|
|
if not self._clusters_l2_normalized():
|
||
|
|
cluster_centers = nn_impl.l2_normalize(cluster_centers, dim=1)
|
||
|
|
|
||
|
|
all_scores, scores, cluster_idx = self._infer_graph(inputs, cluster_centers)
|
||
|
|
if self._use_mini_batch:
|
||
|
|
sync_updates_op = self._mini_batch_sync_updates_op(
|
||
|
|
update_in_steps, cluster_centers_var, cluster_centers_updated,
|
||
|
|
total_counts)
|
||
|
|
assert sync_updates_op is not None
|
||
|
|
with ops.control_dependencies([sync_updates_op]):
|
||
|
|
training_op = self._mini_batch_training_op(inputs, cluster_idx,
|
||
|
|
cluster_centers_updated,
|
||
|
|
total_counts)
|
||
|
|
else:
|
||
|
|
assert cluster_centers == cluster_centers_var
|
||
|
|
training_op = self._full_batch_training_op(inputs, num_clusters,
|
||
|
|
cluster_idx,
|
||
|
|
cluster_centers_var)
|
||
|
|
|
||
|
|
return (all_scores, cluster_idx, scores, cluster_centers_initialized,
|
||
|
|
init_op, training_op)
|
||
|
|
|
||
|
|
def _mini_batch_sync_updates_op(self, update_in_steps, cluster_centers_var,
|
||
|
|
cluster_centers_updated, total_counts):
|
||
|
|
if self._use_mini_batch and self._mini_batch_steps_per_iteration > 1:
|
||
|
|
assert update_in_steps is not None
|
||
|
|
with ops.colocate_with(update_in_steps, ignore_existing=True):
|
||
|
|
|
||
|
|
def _f():
|
||
|
|
# Note that there is a race condition here, so we do a best effort
|
||
|
|
# updates here. We reset update_in_steps first so that other workers
|
||
|
|
# don't duplicate the updates. Also we update cluster_center_vars
|
||
|
|
# before resetting total_counts to avoid large updates to
|
||
|
|
# cluster_centers_updated based on partially updated
|
||
|
|
# cluster_center_vars.
|
||
|
|
with ops.control_dependencies([
|
||
|
|
state_ops.assign(update_in_steps,
|
||
|
|
self._mini_batch_steps_per_iteration - 1)
|
||
|
|
]):
|
||
|
|
with ops.colocate_with(
|
||
|
|
cluster_centers_updated, ignore_existing=True):
|
||
|
|
if self._distance_metric == COSINE_DISTANCE:
|
||
|
|
cluster_centers = nn_impl.l2_normalize(
|
||
|
|
cluster_centers_updated, dim=1)
|
||
|
|
else:
|
||
|
|
cluster_centers = cluster_centers_updated
|
||
|
|
with ops.colocate_with(cluster_centers_var, ignore_existing=True):
|
||
|
|
with ops.control_dependencies(
|
||
|
|
[state_ops.assign(cluster_centers_var, cluster_centers)]):
|
||
|
|
with ops.colocate_with(None, ignore_existing=True):
|
||
|
|
with ops.control_dependencies([
|
||
|
|
state_ops.assign(total_counts,
|
||
|
|
array_ops.zeros_like(total_counts))
|
||
|
|
]):
|
||
|
|
return array_ops.identity(update_in_steps)
|
||
|
|
|
||
|
|
return cond.cond(
|
||
|
|
update_in_steps <= 0, _f,
|
||
|
|
lambda: state_ops.assign_sub(update_in_steps, 1))
|
||
|
|
else:
|
||
|
|
return control_flow_ops.no_op()
|
||
|
|
|
||
|
|
def _mini_batch_training_op(self, inputs, cluster_idx_list, cluster_centers,
|
||
|
|
total_counts):
|
||
|
|
"""Creates an op for training for mini batch case.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
inputs: list of input Tensors.
|
||
|
|
cluster_idx_list: A vector (or list of vectors). Each element in the
|
||
|
|
vector corresponds to an input row in 'inp' and specifies the cluster id
|
||
|
|
corresponding to the input.
|
||
|
|
cluster_centers: Tensor Ref of cluster centers.
|
||
|
|
total_counts: Tensor Ref of cluster counts.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
An op for doing an update of mini-batch k-means.
|
||
|
|
"""
|
||
|
|
update_ops = []
|
||
|
|
for inp, cluster_idx in zip(inputs, cluster_idx_list):
|
||
|
|
with ops.colocate_with(inp, ignore_existing=True):
|
||
|
|
assert total_counts is not None
|
||
|
|
cluster_idx = array_ops.reshape(cluster_idx, [-1])
|
||
|
|
# Dedupe the unique ids of cluster_centers being updated so that updates
|
||
|
|
# can be locally aggregated.
|
||
|
|
unique_ids, unique_idx = array_ops.unique(cluster_idx)
|
||
|
|
num_unique_cluster_idx = array_ops.size(unique_ids)
|
||
|
|
# Fetch the old values of counts and cluster_centers.
|
||
|
|
with ops.colocate_with(total_counts, ignore_existing=True):
|
||
|
|
old_counts = array_ops.gather(total_counts, unique_ids)
|
||
|
|
# TODO(agarwal): This colocation seems to run into problems. Fix it.
|
||
|
|
with ops.colocate_with(cluster_centers, ignore_existing=True):
|
||
|
|
old_cluster_centers = array_ops.gather(cluster_centers, unique_ids)
|
||
|
|
# Locally aggregate the increment to counts.
|
||
|
|
count_updates = math_ops.unsorted_segment_sum(
|
||
|
|
array_ops.ones_like(unique_idx, dtype=total_counts.dtype),
|
||
|
|
unique_idx, num_unique_cluster_idx)
|
||
|
|
# Locally compute the sum of inputs mapped to each id.
|
||
|
|
# For a cluster with old cluster value x, old count n, and with data
|
||
|
|
# d_1,...d_k newly assigned to it, we recompute the new value as
|
||
|
|
# \\(x += (sum_i(d_i) - k * x) / (n + k)\\).
|
||
|
|
# Compute \\(sum_i(d_i)\\), see comment above.
|
||
|
|
cluster_center_updates = math_ops.unsorted_segment_sum(
|
||
|
|
inp, unique_idx, num_unique_cluster_idx)
|
||
|
|
# Shape to enable broadcasting count_updates and learning_rate to inp.
|
||
|
|
# It extends the shape with 1's to match the rank of inp.
|
||
|
|
broadcast_shape = array_ops.concat([
|
||
|
|
array_ops.reshape(num_unique_cluster_idx, [1]),
|
||
|
|
array_ops.ones(
|
||
|
|
array_ops.reshape(array_ops.rank(inp) - 1, [1]),
|
||
|
|
dtype=dtypes.int32)
|
||
|
|
], 0)
|
||
|
|
# Subtract k * x, see comment above.
|
||
|
|
cluster_center_updates -= math_ops.cast(
|
||
|
|
array_ops.reshape(count_updates, broadcast_shape),
|
||
|
|
inp.dtype) * old_cluster_centers
|
||
|
|
learning_rate = math_ops.reciprocal(
|
||
|
|
math_ops.cast(old_counts + count_updates, inp.dtype))
|
||
|
|
learning_rate = array_ops.reshape(learning_rate, broadcast_shape)
|
||
|
|
# scale by 1 / (n + k), see comment above.
|
||
|
|
cluster_center_updates *= learning_rate
|
||
|
|
# Apply the updates.
|
||
|
|
update_counts = state_ops.scatter_add(total_counts, unique_ids,
|
||
|
|
count_updates)
|
||
|
|
update_cluster_centers = state_ops.scatter_add(cluster_centers,
|
||
|
|
unique_ids,
|
||
|
|
cluster_center_updates)
|
||
|
|
update_ops.extend([update_counts, update_cluster_centers])
|
||
|
|
return control_flow_ops.group(*update_ops)
|
||
|
|
|
||
|
|
def _full_batch_training_op(self, inputs, num_clusters, cluster_idx_list,
|
||
|
|
cluster_centers):
|
||
|
|
"""Creates an op for training for full batch case.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
inputs: list of input Tensors.
|
||
|
|
num_clusters: an integer Tensor providing the number of clusters.
|
||
|
|
cluster_idx_list: A vector (or list of vectors). Each element in the
|
||
|
|
vector corresponds to an input row in 'inp' and specifies the cluster id
|
||
|
|
corresponding to the input.
|
||
|
|
cluster_centers: Tensor Ref of cluster centers.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
An op for doing an update of mini-batch k-means.
|
||
|
|
"""
|
||
|
|
cluster_sums = []
|
||
|
|
cluster_counts = []
|
||
|
|
epsilon = constant_op.constant(1e-6, dtype=inputs[0].dtype)
|
||
|
|
for inp, cluster_idx in zip(inputs, cluster_idx_list):
|
||
|
|
with ops.colocate_with(inp, ignore_existing=True):
|
||
|
|
cluster_sums.append(
|
||
|
|
math_ops.unsorted_segment_sum(inp, cluster_idx, num_clusters))
|
||
|
|
cluster_counts.append(
|
||
|
|
math_ops.unsorted_segment_sum(
|
||
|
|
array_ops.reshape(
|
||
|
|
array_ops.ones(
|
||
|
|
array_ops.reshape(array_ops.shape(inp)[0], [-1])),
|
||
|
|
[-1, 1]), cluster_idx, num_clusters))
|
||
|
|
with ops.colocate_with(cluster_centers, ignore_existing=True):
|
||
|
|
new_clusters_centers = math_ops.add_n(cluster_sums) / (
|
||
|
|
math_ops.cast(math_ops.add_n(cluster_counts), cluster_sums[0].dtype) +
|
||
|
|
epsilon)
|
||
|
|
if self._clusters_l2_normalized():
|
||
|
|
new_clusters_centers = nn_impl.l2_normalize(new_clusters_centers, dim=1)
|
||
|
|
return state_ops.assign(cluster_centers, new_clusters_centers)
|
||
|
|
|
||
|
|
|
||
|
|
class _InitializeClustersOpFactory:
|
||
|
|
"""Internal class to create the op to initialize the clusters.
|
||
|
|
|
||
|
|
The op performs this algorithm (see constructor args):
|
||
|
|
|
||
|
|
num_remaining = num_clusters - length(cluster_centers)
|
||
|
|
if num_remaining == 0:
|
||
|
|
assert that cluster_centers_initialized is true
|
||
|
|
else:
|
||
|
|
assert that num_remaining > 0
|
||
|
|
new_centers = choose up to num_remaining initial centers
|
||
|
|
l2-normalize new_centers if using cosine distance
|
||
|
|
all_centers = concat(cluster_centers, new_centers)
|
||
|
|
cluster_centers := all_centers
|
||
|
|
if there is a cluster_centers_updated variable:
|
||
|
|
cluster_centers_updated := cluster_centers
|
||
|
|
num_now_remaining = num_clusters - length(cluster_centers)
|
||
|
|
if num_now_remaining == 0:
|
||
|
|
cluster_centers_initialized := true
|
||
|
|
"""
|
||
|
|
|
||
|
|
# TODO(ccolby): Refactor this class so that kmc2 isn't so much a special case.
|
||
|
|
|
||
|
|
def __init__(self, inputs, num_clusters, initial_clusters, distance_metric,
|
||
|
|
random_seed, kmeans_plus_plus_num_retries, kmc2_chain_length,
|
||
|
|
cluster_centers, cluster_centers_updated,
|
||
|
|
cluster_centers_initialized):
|
||
|
|
"""Creates an op factory.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
inputs: See KMeans constructor.
|
||
|
|
num_clusters: An integer Tensor providing the number of clusters.
|
||
|
|
initial_clusters: See KMeans constructor.
|
||
|
|
distance_metric: See KMeans constructor.
|
||
|
|
random_seed: See KMeans constructor.
|
||
|
|
kmeans_plus_plus_num_retries: See KMeans constructor.
|
||
|
|
kmc2_chain_length: See KMeans constructor.
|
||
|
|
cluster_centers: The TF variable holding the initial centers. It may
|
||
|
|
already contain some centers when the op is executed.
|
||
|
|
cluster_centers_updated: A second TF variable to hold a copy of the
|
||
|
|
initial centers, used for full-batch mode. In mini-batch mode,
|
||
|
|
cluster_centers_updated is the same variable as cluster_centers.
|
||
|
|
cluster_centers_initialized: A boolean TF variable that will be set to
|
||
|
|
true when all the initial centers have been chosen.
|
||
|
|
"""
|
||
|
|
# All of these instance variables are constants.
|
||
|
|
self._inputs = inputs
|
||
|
|
self._num_clusters = num_clusters
|
||
|
|
self._initial_clusters = initial_clusters
|
||
|
|
self._distance_metric = distance_metric
|
||
|
|
self._seed = random_seed
|
||
|
|
self._kmeans_plus_plus_num_retries = kmeans_plus_plus_num_retries
|
||
|
|
self._kmc2_chain_length = kmc2_chain_length
|
||
|
|
self._cluster_centers = cluster_centers
|
||
|
|
self._cluster_centers_updated = cluster_centers_updated
|
||
|
|
self._cluster_centers_initialized = cluster_centers_initialized
|
||
|
|
|
||
|
|
self._num_selected = array_ops.shape(self._cluster_centers)[0]
|
||
|
|
self._num_remaining = self._num_clusters - self._num_selected
|
||
|
|
self._num_data = math_ops.add_n(
|
||
|
|
[array_ops.shape(i)[0] for i in self._inputs])
|
||
|
|
|
||
|
|
def _random(self):
|
||
|
|
indices = random_ops.random_uniform(
|
||
|
|
array_ops.reshape(self._num_remaining, [-1]),
|
||
|
|
minval=0,
|
||
|
|
maxval=math_ops.cast(self._num_data, dtypes.int64),
|
||
|
|
seed=self._seed,
|
||
|
|
dtype=dtypes.int64)
|
||
|
|
return embedding_lookup(self._inputs, indices, partition_strategy='div')
|
||
|
|
|
||
|
|
def _kmeans_plus_plus(self):
|
||
|
|
# Points from only the first shard are used for initializing centers.
|
||
|
|
# TODO(ands): Use all points.
|
||
|
|
inp = self._inputs[0]
|
||
|
|
if self._distance_metric == COSINE_DISTANCE:
|
||
|
|
inp = nn_impl.l2_normalize(inp, dim=1)
|
||
|
|
return gen_clustering_ops.kmeans_plus_plus_initialization(
|
||
|
|
inp, math_ops.cast(self._num_remaining, dtypes.int64), self._seed,
|
||
|
|
self._kmeans_plus_plus_num_retries)
|
||
|
|
|
||
|
|
def _kmc2_multiple_centers(self):
|
||
|
|
"""Adds new initial cluster centers using the k-MC2 algorithm.
|
||
|
|
|
||
|
|
In each call to the op, the provided batch is split into subsets based on
|
||
|
|
the specified `kmc2_chain_length`. On each subset, a single Markov chain of
|
||
|
|
the k-MC2 algorithm is used to add *one* new center cluster center. If there
|
||
|
|
are less than `kmc2_chain_length` points in the subset, a single center is
|
||
|
|
added using one Markov chain on the full input. It is assumed that the
|
||
|
|
provided batch has previously been randomly permuted. Otherwise, k-MC2 may
|
||
|
|
return suboptimal centers.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
An op that adds new cluster centers.
|
||
|
|
"""
|
||
|
|
# The op only operates on the first shard of data.
|
||
|
|
first_shard = self._inputs[0]
|
||
|
|
# Number of points in the input that can be used.
|
||
|
|
batch_size = array_ops.shape(first_shard)[0]
|
||
|
|
# Maximum number of subsets such that the size of each subset is at least
|
||
|
|
# `kmc2_chain_length`. Final subsets may be larger.
|
||
|
|
max_to_sample = math_ops.cast(
|
||
|
|
batch_size / self._kmc2_chain_length, dtype=dtypes.int32)
|
||
|
|
# We sample at least one new center and at most all remaining centers.
|
||
|
|
num_to_sample = math_ops.maximum(
|
||
|
|
math_ops.minimum(self._num_remaining, max_to_sample), 1)
|
||
|
|
|
||
|
|
def _cond(i, _):
|
||
|
|
"""Stopping condition for the while loop."""
|
||
|
|
return math_ops.less(i, num_to_sample)
|
||
|
|
|
||
|
|
def _body(i, _):
|
||
|
|
"""Body that adds a single new center based on a subset."""
|
||
|
|
|
||
|
|
def _sample_random():
|
||
|
|
"""Returns a random point as a cluster center."""
|
||
|
|
# By assumption the batch is reshuffled and _sample_random is always
|
||
|
|
# called for i=0. Hence, we simply return the first point.
|
||
|
|
new_center = array_ops.reshape(first_shard[0], [1, -1])
|
||
|
|
if self._distance_metric == COSINE_DISTANCE:
|
||
|
|
new_center = nn_impl.l2_normalize(new_center, dim=1)
|
||
|
|
return new_center
|
||
|
|
|
||
|
|
def _sample_kmc2_chain():
|
||
|
|
"""Returns previous centers as well as a new center sampled using k-MC2."""
|
||
|
|
# Extract the subset from the underlying batch.
|
||
|
|
start = i * self._kmc2_chain_length
|
||
|
|
end = start + self._kmc2_chain_length
|
||
|
|
subset = first_shard[start:end]
|
||
|
|
# Compute the distances from points in the subset to previous centers.
|
||
|
|
_, distances = gen_clustering_ops.nearest_neighbors(
|
||
|
|
subset, self._cluster_centers, 1)
|
||
|
|
# Sample index of new center using k-MC2 Markov chain.
|
||
|
|
new_center_index = gen_clustering_ops.kmc2_chain_initialization(
|
||
|
|
array_ops.squeeze(distances), self._seed)
|
||
|
|
# Extract actual new center.
|
||
|
|
newly_sampled_center = array_ops.reshape(subset[new_center_index],
|
||
|
|
[1, -1])
|
||
|
|
# Return concatenation with previously sampled centers.
|
||
|
|
if self._distance_metric == COSINE_DISTANCE:
|
||
|
|
newly_sampled_center = nn_impl.l2_normalize(
|
||
|
|
newly_sampled_center, dim=1)
|
||
|
|
return array_ops.concat([self._cluster_centers, newly_sampled_center],
|
||
|
|
0)
|
||
|
|
|
||
|
|
# Obtain a random point if there are no previously sampled centers.
|
||
|
|
# Otherwise, construct a k-MC2 Markov chain.
|
||
|
|
new_centers = cond.cond(
|
||
|
|
math_ops.equal(self._num_selected, 0), _sample_random,
|
||
|
|
_sample_kmc2_chain)
|
||
|
|
# Assign new cluster centers to underlying variable.
|
||
|
|
assigned_centers = state_ops.assign(
|
||
|
|
self._cluster_centers, new_centers, validate_shape=False)
|
||
|
|
if self._cluster_centers_updated is not self._cluster_centers:
|
||
|
|
assigned_centers = state_ops.assign(
|
||
|
|
self._cluster_centers_updated,
|
||
|
|
assigned_centers,
|
||
|
|
validate_shape=False)
|
||
|
|
return i + 1, self._num_clusters - array_ops.shape(assigned_centers)[0]
|
||
|
|
|
||
|
|
# Add num_to_sample new data points.
|
||
|
|
_, num_remaining = while_loop.while_loop(_cond, _body, [0, 0])
|
||
|
|
return num_remaining
|
||
|
|
|
||
|
|
def _greedy_batch_sampler(self, sampler):
|
||
|
|
# If the input dataset size is smaller than the number of centers
|
||
|
|
# remaining, choose the entire input dataset as centers. This can happen
|
||
|
|
# with mini-batch. Otherwise, sample the batch according to the provided
|
||
|
|
# sampler.
|
||
|
|
return cond.cond(self._num_data <= self._num_remaining,
|
||
|
|
lambda: array_ops.concat(self._inputs, 0),
|
||
|
|
sampler)
|
||
|
|
|
||
|
|
def _single_batch_sampler(self, sampler):
|
||
|
|
# Enforce that there are at least as many data points as centers
|
||
|
|
# remaining. This gives the provided sampler the chance to select all
|
||
|
|
# remaining centers from a single batch.
|
||
|
|
with ops.control_dependencies(
|
||
|
|
[check_ops.assert_greater_equal(self._num_data, self._num_remaining)]):
|
||
|
|
return sampler()
|
||
|
|
|
||
|
|
def _choose_initial_centers(self):
|
||
|
|
if isinstance(self._initial_clusters, str):
|
||
|
|
if self._initial_clusters == RANDOM_INIT:
|
||
|
|
return self._greedy_batch_sampler(self._random)
|
||
|
|
else: # self._initial_clusters == KMEANS_PLUS_PLUS_INIT
|
||
|
|
return self._single_batch_sampler(self._kmeans_plus_plus)
|
||
|
|
elif callable(self._initial_clusters):
|
||
|
|
return self._initial_clusters(self._inputs, self._num_remaining)
|
||
|
|
else:
|
||
|
|
with ops.control_dependencies([
|
||
|
|
check_ops.assert_equal(self._num_remaining,
|
||
|
|
array_ops.shape(self._initial_clusters)[0])
|
||
|
|
]):
|
||
|
|
return self._initial_clusters
|
||
|
|
|
||
|
|
def _add_new_centers(self):
|
||
|
|
"""Adds some centers and returns the number of centers remaining."""
|
||
|
|
new_centers = self._choose_initial_centers()
|
||
|
|
if self._distance_metric == COSINE_DISTANCE:
|
||
|
|
new_centers = nn_impl.l2_normalize(new_centers, dim=1)
|
||
|
|
# If cluster_centers is empty, it doesn't have the right shape for concat.
|
||
|
|
all_centers = cond.cond(
|
||
|
|
math_ops.equal(self._num_selected, 0), lambda: new_centers,
|
||
|
|
lambda: array_ops.concat([self._cluster_centers, new_centers], 0))
|
||
|
|
# TODO(ccolby): De-dupe all_centers?
|
||
|
|
a = state_ops.assign(
|
||
|
|
self._cluster_centers, all_centers, validate_shape=False)
|
||
|
|
if self._cluster_centers_updated is not self._cluster_centers:
|
||
|
|
a = state_ops.assign(
|
||
|
|
self._cluster_centers_updated, a, validate_shape=False)
|
||
|
|
return self._num_clusters - array_ops.shape(a)[0]
|
||
|
|
|
||
|
|
def _initialize(self):
|
||
|
|
with ops.control_dependencies([
|
||
|
|
check_ops.assert_positive(self._num_remaining),
|
||
|
|
]):
|
||
|
|
if self._initial_clusters == KMC2_INIT:
|
||
|
|
num_now_remaining = self._kmc2_multiple_centers()
|
||
|
|
else:
|
||
|
|
num_now_remaining = self._add_new_centers()
|
||
|
|
return cond.cond(
|
||
|
|
math_ops.equal(num_now_remaining, 0),
|
||
|
|
lambda: state_ops.assign(self._cluster_centers_initialized, True),
|
||
|
|
control_flow_ops.no_op)
|
||
|
|
|
||
|
|
def op(self):
|
||
|
|
"""Returns the cluster initializer op."""
|
||
|
|
return cond.cond(
|
||
|
|
math_ops.equal(self._num_remaining, 0),
|
||
|
|
lambda: check_ops.assert_equal(self._cluster_centers_initialized, True),
|
||
|
|
self._initialize)
|