Moving kmeans from cuml to Raft

cpp/include/raft/cluster/kmeans.cuh

@@ -15,51 +15,129 @@
  */
 #pragma once
 
+#include <optional>
 #include <raft/cluster/detail/kmeans.cuh>
+#include <raft/cluster/kmeans_params.hpp>
+#include <raft/core/mdarray.hpp>
 
 namespace raft {
 namespace cluster {
 
 /**
- *  @brief Find clusters with k-means algorithm.
- *    Initial centroids are chosen with k-means++ algorithm. Empty
- *    clusters are reinitialized by choosing new centroids with
- *    k-means++ algorithm.
- *  @tparam index_type_t the type of data used for indexing.
- *  @tparam value_type_t the type of data used for weights, distances.
- *  @param handle the raft handle.
- *  @param n Number of observation vectors.
- *  @param d Dimension of observation vectors.
- *  @param k Number of clusters.
- *  @param tol Tolerance for convergence. k-means stops when the
- *    change in residual divided by n is less than tol.
- *  @param maxiter Maximum number of k-means iterations.
- *  @param obs (Input, device memory, d*n entries) Observation
- *    matrix. Matrix is stored column-major and each column is an
- *    observation vector. Matrix dimensions are d x n.
- *  @param codes (Output, device memory, n entries) Cluster
- *    assignments.
- *  @param residual On exit, residual sum of squares (sum of squares
- *    of distances between observation vectors and centroids).
- *  @param iters on exit, number of k-means iterations.
- *  @param seed random seed to be used.
- *  @return error flag
+ * @brief Find clusters with k-means algorithm.
+ *   Initial centroids are chosen with k-means++ algorithm. Empty
+ *   clusters are reinitialized by choosing new centroids with
+ *   k-means++ algorithm.
+ * @tparam DataT the type of data used for weights, distances.
+ * @tparam IdxT the type of data used for indexing.
+ * @param[in]     handle        The raft handle.
+ * @param[in]     params        Parameters for KMeans model.
+ * @param[in]     X             Training instances to cluster. The data must
+ * be in row-major format
+ * @param[in]     sample_weight Optional weights for each observation in X.
+ * @param[inout]  centroids     [in] When init is InitMethod::Array, use
+ * centroids as the initial cluster centers
+ *                              [out] Otherwise, generated centroids from the
+ * kmeans algorithm is stored at the address pointed by 'centroids'.
+ * @param[out]    inertia       Sum of squared distances of samples to their
+ * closest cluster center.
+ * @param[out]    n_iter        Number of iterations run.
  */
-template <typename index_type_t, typename value_type_t>
-int kmeans(handle_t const& handle,
-           index_type_t n,
-           index_type_t d,
-           index_type_t k,
-           value_type_t tol,
-           index_type_t maxiter,
-           const value_type_t* __restrict__ obs,
-           index_type_t* __restrict__ codes,
-           value_type_t& residual,
-           index_type_t& iters,
-           unsigned long long seed = 123456)
+template <typename DataT, typename IndexT = int>
+void kmeans_fit(handle_t const& handle,
+                const KMeansParams& params,
+                raft::device_matrix_view<const DataT> X,
+                std::optional<raft::device_vector_view<const DataT>> sample_weight,
+                raft::device_matrix_view<DataT> centroids,
+                DataT& inertia,
+                IndexT& n_iter)
 {
-  return detail::kmeans<index_type_t, value_type_t>(
-    handle, n, d, k, tol, maxiter, obs, codes, residual, iters, seed);
+  detail::kmeans_fit<DataT, IndexT>(handle, params, X, sample_weight, centroids, inertia, n_iter);
+}
+
+/**
+ * @brief Predict the closest cluster each sample in X belongs to.
+ *
+ * @param[in]     handle            The handle to the cuML library context
+ * that manages the CUDA resources.
+ * @param[in]     params            Parameters for KMeans model.
+ * @param[in]     centroids         Cluster centroids. The data must be in
+ * row-major format.
+ * @param[in]     X                 New data to predict.
+ * @param[in]     sample_weight     The weights for each observation in X.
+ * @param[in]     normalize_weight  True if the weights should be normalized
+ * @param[out]    labels            Index of the cluster each sample in X
+ * belongs to.
+ * @param[out]    inertia           Sum of squared distances of samples to
+ * their closest cluster center.
+ */
+template <typename DataT, typename IndexT = int>
+void kmeans_predict(handle_t const& handle,
+                    const KMeansParams& params,
+                    raft::device_matrix_view<const DataT> X,
+                    std::optional<raft::device_vector_view<const DataT>> sample_weight,
+                    raft::device_matrix_view<const DataT> centroids,
+                    raft::device_vector_view<IndexT> labels,
+                    bool normalize_weight,
+                    DataT& inertia)
+{
+  detail::kmeans_predict<DataT, IndexT>(
+    handle, params, X, sample_weight, centroids, labels, normalize_weight, inertia);
+}
+
+/**
+ * @brief Compute k-means clustering and predicts cluster index for each sample
+ * in the input.
+ *
+ * @param[in]     handle        The handle to the cuML library context that
+ * manages the CUDA resources.
+ * @param[in]     params        Parameters for KMeans model.
+ * @param[in]     X             Training instances to cluster. The data must be
+ * in row-major format
+ * @param[in]     sample_weight The weights for each observation in X.
+ * @param[inout]  centroids     [in] When init is InitMethod::Array, use
+ * centroids  as the initial cluster centers
+ *                              [out] Otherwise, generated centroids from the
+ * kmeans algorithm is stored at the address pointed by 'centroids'.
+ * @param[out]    labels        Index of the cluster each sample in X belongs
+ * to.
+ * @param[out]    inertia       Sum of squared distances of samples to their
+ * closest cluster center.
+ * @param[out]    n_iter        Number of iterations run.
+ */
+template <typename DataT, typename IndexT = int>
+void kmeans_fit_predict(handle_t const& handle,
+                        const KMeansParams& params,
+                        raft::device_matrix_view<const DataT> X,
+                        std::optional<raft::device_vector_view<const DataT>> sample_weight,
+                        raft::device_matrix_view<DataT> centroids,
+                        raft::device_vector_view<IndexT> labels,
+                        DataT& inertia,
+                        IndexT& n_iter)
+{
+  kmeans_fit<DataT, IndexT>(handle, params, X, sample_weight, centroids, inertia, n_iter);
+  kmeans_predict<DataT, IndexT>(handle, params, X, sample_weight, centroids, labels, true, inertia);
+}
+
+/**
+ * @brief Transform X to a cluster-distance space.
+ *
+ * @param[in]     handle        The handle to the cuML library context that
+ * manages the CUDA resources.
+ * @param[in]     params        Parameters for KMeans model.
+ * @param[in]     X             Training instances to cluster. The data must
+ * be in row-major format
+ * @param[in]     centroids     Cluster centroids. The data must be in row-major format.
+ * @param[out]    X_new         X transformed in the new space..
+ */
+template <typename DataT, typename IndexT = int>
+void kmeans_transform(const raft::handle_t& handle,
+                      const KMeansParams& params,
+                      raft::device_matrix_view<const DataT> X,
+                      raft::device_matrix_view<const DataT> centroids,
+                      raft::device_matrix_view<DataT> X_new)
+{
+  detail::kmeans_transform<DataT, IndexT>(handle, params, X, centroids, X_new);
 }
 }  // namespace cluster
 }  // namespace raft

cpp/include/raft/cluster/kmeans_params.hpp

@@ -0,0 +1,71 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.
+ *
+ * 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.
+ */
+#pragma once
+//#include <raft/common/logger.hpp>
+
+namespace raft {
+namespace cluster {
+
+struct KMeansParams {
+  enum InitMethod { KMeansPlusPlus, Random, Array };
+
+  // The number of clusters to form as well as the number of centroids to
+  // generate (default:8).
+  int n_clusters = 8;
+
+  /*
+   * Method for initialization, defaults to k-means++:
+   *  - InitMethod::KMeansPlusPlus (k-means++): Use scalable k-means++ algorithm
+   * to select the initial cluster centers.
+   *  - InitMethod::Random (random): Choose 'n_clusters' observations (rows) at
+   * random from the input data for the initial centroids.
+   *  - InitMethod::Array (ndarray): Use 'centroids' as initial cluster centers.
+   */
+  InitMethod init = KMeansPlusPlus;
+
+  // Maximum number of iterations of the k-means algorithm for a single run.
+  int max_iter = 300;
+
+  // Relative tolerance with regards to inertia to declare convergence.
+  double tol = 1e-4;
+
+  // verbosity level.
+  int verbosity = 4;  // RAFT_LEVEL_INFO;
+
+  // Seed to the random number generator.
+  uint64_t seed = 0;
+
+  // Metric to use for distance computation. Any metric from
+  // raft::distance::DistanceType can be used
+  int metric = 0;
+
+  // Number of instance k-means algorithm will be run with different seeds.
+  int n_init = 1;
+
+  // Oversampling factor for use in the k-means|| algorithm.
+  double oversampling_factor = 2.0;
+
+  // batch_samples and batch_centroids are used to tile 1NN computation which is
+  // useful to optimize/control the memory footprint
+  // Default tile is [batch_samples x n_clusters] i.e. when batch_centroids is 0
+  // then don't tile the centroids
+  int batch_samples   = 1 << 15;
+  int batch_centroids = 0;  // if 0 then batch_centroids = n_clusters
+
+  bool inertia_check = false;
+};
+}  // namespace cluster
+}  // namespace raft