TIGER / GRID /src /components /clustering_initializers.py

Upload TIGER project for AutoDL

d9115cb verified 8 months ago

9.35 kB

from abc import abstractmethod
from functools import partial

import torch
import torch.nn as nn
from pytorch_lightning import LightningModule

from src.components.distance_functions import DistanceFunction


class ClusteringInitializer(nn.Module):
    """
    Base class for clustering initializers.
    This class provides an interface for initializing centroids for clustering algorithms.
    """

    def __init__(self, n_clusters: int, initialize_on_cpu: bool = False):
        """
        Initializes the ClusteringInitializer class.

        Args:
            n_clusters: Number of clusters to form.
            initialize_on_cpu: Whether to move the tensors to the CPU for computing the
                initialization. This is useful for large initialization buffer sizes for
                which GPU memory might be a constraint. Otherwise, it is recommended to
                keep the tensors on the GPU for faster computation.
        """
        super().__init__()
        self.n_clusters = n_clusters
        self.initialize_on_cpu = initialize_on_cpu

    @abstractmethod
    def forward(self, buffer: torch.Tensor) -> torch.Tensor:
        """
        Initialize centroids for clustering algorithms.
        Args:
            buffer: Data points of shape (batch_size, n_features)
        Returns:
            Initialized centroids of shape (n_clusters, n_features)
        """
        pass


class RandomInitializer(ClusteringInitializer):
    """
    Random initialization for clustering algorithms.
    This class provides a method to initialize centroids randomly from the data points.
    """

    def __init__(self, n_clusters: int, initialize_on_cpu: bool = True):
        """
        Initializes the RandomInitialize class with the specified number of clusters.
        Args:
            n_clusters: Number of clusters to form.
        """
        super().__init__(n_clusters=n_clusters, initialize_on_cpu=initialize_on_cpu)

    def forward(self, buffer: torch.Tensor) -> torch.Tensor:
        """
        Initialize centroids randomly from the data points.
        Args:
            buffer: Data points of shape (batch_size, n_features)
        Returns:
            Initialized centroids of shape (n_clusters, n_features)
        """
        if self.initialize_on_cpu:
            old_device = buffer.device
            # Move the buffer to CPU for initialization
            buffer = buffer.to("cpu")

        n_samples = buffer.shape[0]

        # Choose centroids randomly from the data points
        indices = torch.randperm(n_samples, device=buffer.device)[: self.n_clusters]
        centroids = buffer[indices].clone().data
        if self.initialize_on_cpu:
            # Move the centroids back to the original device
            centroids = centroids.to(old_device)

        return centroids


class KMeansPlusPlusInitInitializer(ClusteringInitializer):
    """
    KMeans++ initialization for clustering algorithms.
    This class provides a method to initialize centroids using the k-means++
    algorithm, which is a smarter way to choose the initial centroids for k-means
    clustering. It helps in faster convergence and better clustering results.
    """

    def __init__(
        self,
        n_clusters: int,
        distance_function: DistanceFunction,
        initialize_on_cpu: bool = True,
    ):
        """
        Initializes the KMeansPlusPlusInitInitializer class with the specified parameters.
        Args:
            n_clusters: Number of clusters to form.
            distance_function: Function to compute distances between points.
            initialize_on_cpu: Whether to move the tensors to the CPU for computing the
                initialization.
        """
        super().__init__(n_clusters=n_clusters, initialize_on_cpu=initialize_on_cpu)
        self.distance_function = distance_function

    def forward(self, buffer: torch.Tensor) -> torch.Tensor:
        """
        Initialize centroids using k-means++ algorithm for better convergence.
        The k-means++ algorithm iteratively samples the next initial centroid by
        sampling points with probability proportional to their distance to the nearest
        existing centroid.
        Args:
            buffer: Data points of shape (batch_size, n_features)
        Returns:
            Initialized centroids of shape (n_clusters, n_features)
        """
        if self.initialize_on_cpu:
            old_device = buffer.device
            # Move the buffer to CPU for initialization
            buffer = buffer.to("cpu")

        n_samples = buffer.shape[0]
        n_features = buffer.shape[1]
        centroids = torch.zeros(
            (self.n_clusters, n_features), dtype=buffer.dtype, device=buffer.device
        )

        # Choose first centroid randomly
        first_centroid_idx = torch.randint(0, n_samples, (1,), device=buffer.device)
        centroids[0] = buffer[first_centroid_idx]

        # Choose remaining centroids
        # We cannot vectorize this part because it is inherently sequential
        # However, we only need to execute this loop once at initialization
        for i in range(1, self.n_clusters):
            # Compute distances to the nearest existing centroid
            min_distances = torch.min(
                self.distance_function.compute(buffer, centroids[:i]), dim=1
            )[0]
            if min_distances.sum() == 0:
                # All points are already centroids, so we simply assign the remaining
                # centroids randomly
                centroids[i:] = buffer[
                    torch.randint(
                        0, n_samples, (self.n_clusters - i,), device=buffer.device
                    )
                ]
                break

            # Choose the next centroid with probability proportional to distance
            next_centroid_idx = torch.multinomial(min_distances, num_samples=1)

            # Assign the next centroid
            # We do not need to remove the point from buffer because the probability
            # will be zero for all future iterations
            centroids[i] = buffer[next_centroid_idx]

        if self.initialize_on_cpu:
            # Move the centroids back to the original device
            centroids = centroids.to(old_device)

        return centroids


class ClusteringModuleInitializer(ClusteringInitializer):
    """
    Module to initialize clustering algorithm with the result of another
    clustering algorithm.
    """

    def __init__(
        self,
        n_clusters: int,
        clustering_module: LightningModule,
        initialize_on_cpu: bool = False,
        max_iter: int = 100,
        atol: float = 1e-8,
    ):
        """
        Initializes the KMeansInitializer class with the specified parameters.

        Args:
            n_clusters: Number of clusters to form.
            clustering_module: An instance of a `BaseClusteringModule` to be used to
                initialize the centroids.
            initialize_on_cpu: Whether to move the tensors to the CPU for computing the
                initialization. This is useful for large initialization buffer sizes for
                which GPU memory might be a constraint. Otherwise, it is recommended to
                keep the tensors on the GPU for faster computation.
            max_iter: Maximum number of iterations to run the clustering module for.
            atol: Absolute tolerance for convergence. If all elements of the centroids
                do not change more than this value on consecutive iterations, the
                initialization is considered converged.
        """
        super().__init__(n_clusters=n_clusters, initialize_on_cpu=initialize_on_cpu)

        from src.models.modules.clustering.base_clustering_module import (  # we import here to avoid circular imports
            BaseClusteringModule,
        )

        assert isinstance(
            clustering_module, BaseClusteringModule
        ), "clustering_module must be an instance of BaseClusteringModule"

        self.clustering_module = clustering_module
        self.max_iter = max_iter
        self.atol = atol

    def forward(self, buffer: torch.Tensor) -> torch.Tensor:
        """
        Initialize centroids using self.clustering_module.

        Args:
            buffer: Data points of shape (batch_size, n_features)

        Returns:
            Initialized centroids of shape (n_clusters, n_features)
        """
        self.clustering_module.on_train_start()
        cur_centroids = self.clustering_module.get_centroids()
        for step in range(self.max_iter):
            # Perform a training step
            self.clustering_module.model_step(buffer)
            new_centroids = self.clustering_module.get_centroids()

            # Check for convergence
            if step > 0 and torch.allclose(
                cur_centroids, new_centroids, atol=self.atol
            ):
                print(f"Initialization converged after {step} iterations")
                break
            cur_centroids = new_centroids.clone()

        return cur_centroids.detach()

	from abc import abstractmethod
	from functools import partial

	import torch
	import torch.nn as nn
	from pytorch_lightning import LightningModule

	from src.components.distance_functions import DistanceFunction


	class ClusteringInitializer(nn.Module):
	"""
	Base class for clustering initializers.
	This class provides an interface for initializing centroids for clustering algorithms.
	"""

	def __init__(self, n_clusters: int, initialize_on_cpu: bool = False):
	"""
	Initializes the ClusteringInitializer class.

	Args:
	n_clusters: Number of clusters to form.
	initialize_on_cpu: Whether to move the tensors to the CPU for computing the
	initialization. This is useful for large initialization buffer sizes for
	which GPU memory might be a constraint. Otherwise, it is recommended to
	keep the tensors on the GPU for faster computation.
	"""
	super().__init__()
	self.n_clusters = n_clusters
	self.initialize_on_cpu = initialize_on_cpu

	@abstractmethod
	def forward(self, buffer: torch.Tensor) -> torch.Tensor:
	"""
	Initialize centroids for clustering algorithms.
	Args:
	buffer: Data points of shape (batch_size, n_features)
	Returns:
	Initialized centroids of shape (n_clusters, n_features)
	"""
	pass


	class RandomInitializer(ClusteringInitializer):
	"""
	Random initialization for clustering algorithms.
	This class provides a method to initialize centroids randomly from the data points.
	"""

	def __init__(self, n_clusters: int, initialize_on_cpu: bool = True):
	"""
	Initializes the RandomInitialize class with the specified number of clusters.
	Args:
	n_clusters: Number of clusters to form.
	"""
	super().__init__(n_clusters=n_clusters, initialize_on_cpu=initialize_on_cpu)

	def forward(self, buffer: torch.Tensor) -> torch.Tensor:
	"""
	Initialize centroids randomly from the data points.
	Args:
	buffer: Data points of shape (batch_size, n_features)
	Returns:
	Initialized centroids of shape (n_clusters, n_features)
	"""
	if self.initialize_on_cpu:
	old_device = buffer.device
	# Move the buffer to CPU for initialization
	buffer = buffer.to("cpu")

	n_samples = buffer.shape[0]

	# Choose centroids randomly from the data points
	indices = torch.randperm(n_samples, device=buffer.device)[: self.n_clusters]
	centroids = buffer[indices].clone().data
	if self.initialize_on_cpu:
	# Move the centroids back to the original device
	centroids = centroids.to(old_device)

	return centroids


	class KMeansPlusPlusInitInitializer(ClusteringInitializer):
	"""
	KMeans++ initialization for clustering algorithms.
	This class provides a method to initialize centroids using the k-means++
	algorithm, which is a smarter way to choose the initial centroids for k-means
	clustering. It helps in faster convergence and better clustering results.
	"""

	def __init__(
	self,
	n_clusters: int,
	distance_function: DistanceFunction,
	initialize_on_cpu: bool = True,
	):
	"""
	Initializes the KMeansPlusPlusInitInitializer class with the specified parameters.
	Args:
	n_clusters: Number of clusters to form.
	distance_function: Function to compute distances between points.
	initialize_on_cpu: Whether to move the tensors to the CPU for computing the
	initialization.
	"""
	super().__init__(n_clusters=n_clusters, initialize_on_cpu=initialize_on_cpu)
	self.distance_function = distance_function

	def forward(self, buffer: torch.Tensor) -> torch.Tensor:
	"""
	Initialize centroids using k-means++ algorithm for better convergence.
	The k-means++ algorithm iteratively samples the next initial centroid by
	sampling points with probability proportional to their distance to the nearest
	existing centroid.
	Args:
	buffer: Data points of shape (batch_size, n_features)
	Returns:
	Initialized centroids of shape (n_clusters, n_features)
	"""
	if self.initialize_on_cpu:
	old_device = buffer.device
	# Move the buffer to CPU for initialization
	buffer = buffer.to("cpu")

	n_samples = buffer.shape[0]
	n_features = buffer.shape[1]
	centroids = torch.zeros(
	(self.n_clusters, n_features), dtype=buffer.dtype, device=buffer.device
	)

	# Choose first centroid randomly
	first_centroid_idx = torch.randint(0, n_samples, (1,), device=buffer.device)
	centroids[0] = buffer[first_centroid_idx]

	# Choose remaining centroids
	# We cannot vectorize this part because it is inherently sequential
	# However, we only need to execute this loop once at initialization
	for i in range(1, self.n_clusters):
	# Compute distances to the nearest existing centroid
	min_distances = torch.min(
	self.distance_function.compute(buffer, centroids[:i]), dim=1
	)[0]
	if min_distances.sum() == 0:
	# All points are already centroids, so we simply assign the remaining
	# centroids randomly
	centroids[i:] = buffer[
	torch.randint(
	0, n_samples, (self.n_clusters - i,), device=buffer.device
	)
	]
	break

	# Choose the next centroid with probability proportional to distance
	next_centroid_idx = torch.multinomial(min_distances, num_samples=1)

	# Assign the next centroid
	# We do not need to remove the point from buffer because the probability
	# will be zero for all future iterations
	centroids[i] = buffer[next_centroid_idx]

	if self.initialize_on_cpu:
	# Move the centroids back to the original device
	centroids = centroids.to(old_device)

	return centroids


	class ClusteringModuleInitializer(ClusteringInitializer):
	"""
	Module to initialize clustering algorithm with the result of another
	clustering algorithm.
	"""

	def __init__(
	self,
	n_clusters: int,
	clustering_module: LightningModule,
	initialize_on_cpu: bool = False,
	max_iter: int = 100,
	atol: float = 1e-8,
	):
	"""
	Initializes the KMeansInitializer class with the specified parameters.

	Args:
	n_clusters: Number of clusters to form.
	clustering_module: An instance of a `BaseClusteringModule` to be used to
	initialize the centroids.
	initialize_on_cpu: Whether to move the tensors to the CPU for computing the
	initialization. This is useful for large initialization buffer sizes for
	which GPU memory might be a constraint. Otherwise, it is recommended to
	keep the tensors on the GPU for faster computation.
	max_iter: Maximum number of iterations to run the clustering module for.
	atol: Absolute tolerance for convergence. If all elements of the centroids
	do not change more than this value on consecutive iterations, the
	initialization is considered converged.
	"""
	super().__init__(n_clusters=n_clusters, initialize_on_cpu=initialize_on_cpu)

	from src.models.modules.clustering.base_clustering_module import ( # we import here to avoid circular imports
	BaseClusteringModule,
	)

	assert isinstance(
	clustering_module, BaseClusteringModule
	), "clustering_module must be an instance of BaseClusteringModule"

	self.clustering_module = clustering_module
	self.max_iter = max_iter
	self.atol = atol

	def forward(self, buffer: torch.Tensor) -> torch.Tensor:
	"""
	Initialize centroids using self.clustering_module.

	Args:
	buffer: Data points of shape (batch_size, n_features)

	Returns:
	Initialized centroids of shape (n_clusters, n_features)
	"""
	self.clustering_module.on_train_start()
	cur_centroids = self.clustering_module.get_centroids()
	for step in range(self.max_iter):
	# Perform a training step
	self.clustering_module.model_step(buffer)
	new_centroids = self.clustering_module.get_centroids()

	# Check for convergence
	if step > 0 and torch.allclose(
	cur_centroids, new_centroids, atol=self.atol
	):
	print(f"Initialization converged after {step} iterations")
	break
	cur_centroids = new_centroids.clone()

	return cur_centroids.detach()