sampler.py

import abc
from typing import Union

import numpy as np
import torch
import tqdm

import pickle
import os

import utils

class IdentitySampler:
    def run(
        self, features: Union[torch.Tensor, np.ndarray]
    ) -> Union[torch.Tensor, np.ndarray]:
        return features


class BaseSampler(abc.ABC):
    def __init__(self, percentage: float):
        if not 0 < percentage < 1:
            raise ValueError("Percentage value not in (0, 1).")
        self.percentage = percentage

    @abc.abstractmethod
    def run(
        self, features: Union[torch.Tensor, np.ndarray]
    ) -> Union[torch.Tensor, np.ndarray]:
        pass

    def _store_type(self, features: Union[torch.Tensor, np.ndarray]) -> None:
        self.features_is_numpy = isinstance(features, np.ndarray)
        if not self.features_is_numpy:
            self.features_device = features.device

    def _restore_type(self, features: torch.Tensor) -> Union[torch.Tensor, np.ndarray]:
        if self.features_is_numpy:
            return features.cpu().numpy()
        return features.to(self.features_device)


class GreedyCoresetSampler(BaseSampler):
    def __init__(
        self,
        percentage: float,
        device: torch.device,
        seed: int = 0, 
        batch_size: int = -1,
        is_reduce: bool = True,
        method: str = 'random_proj',
        dimension_to_project_features_to=128
    ):
        """Greedy Coreset sampling base class."""
        super().__init__(percentage)

        self.batch_size = batch_size
        self.device = device
        self.dimension_to_project_features_to = dimension_to_project_features_to
        self.is_reduce = is_reduce
        self.method = method
        
        self.seed = seed
        

    def _reduce_features(self, features):
        if features.shape[1] == self.dimension_to_project_features_to:
            return features
        
        if self.method == 'random_proj':
            print('Reducing features using random projection.')
            features = self.random_projection(features)
        elif self.method == 'linear_interpolation':
            print('Reducing features using linear interpolation.')
            features = self.linear_interpolation(features)
        
        return features
        
    def linear_interpolation(self, features):
        return torch.nn.functional.interpolate(
            features.unsqueeze(1), 
            size=self.dimension_to_project_features_to,
            mode='linear'
        ).squeeze(1)
        
            
    def random_projection(self, features):
        utils.fix_seeds(seed=self.seed)
        mapper = torch.nn.Linear(
            features.shape[1], self.dimension_to_project_features_to, bias=False
        )
        mapper.reset_parameters()
        _ = mapper.to(self.device)
        
        with torch.no_grad():
            if self.batch_size != -1:
                reduced_features = []
                row_nb_step = int(np.ceil(len(features) / self.batch_size))
                
                for r_step in tqdm.tqdm(np.arange(row_nb_step), desc='Reducing features dimension...'):
                    row_batch_start = r_step * self.batch_size
                    row_batch_end = min((r_step + 1) * self.batch_size, len(features))
                    features_b = features[row_batch_start: row_batch_end].to(self.device)
                    
                    reduced_features.append(mapper(features_b))
                    
                return torch.cat(reduced_features)
            else:    
                features = features.to(self.device)
                return mapper(features)

    def run(
        self, features: Union[torch.Tensor, np.ndarray]
    ) -> Union[torch.Tensor, np.ndarray]:
        """Subsamples features using Greedy Coreset.
        Args:
            features: [N x D]
        """
        if self.percentage == 1:
            return features
        self._store_type(features)
        if isinstance(features, np.ndarray):
            features = torch.from_numpy(features)
            
        if self.is_reduce:
            reduced_features = self._reduce_features(features)
        else:
            print('No reducing features.')
            reduced_features = features.to(self.device)
        
        with torch.no_grad():
            sample_indices = self._compute_greedy_coreset_indices(reduced_features)
        features = features[sample_indices]
        return self._restore_type(features), sample_indices

    def _compute_greedy_coreset_indices(self, features: torch.Tensor) -> np.ndarray:
        """Runs iterative greedy coreset selection.
        Args:
            features: [NxD] input feature bank to sample.
        """
        if self.batch_size == -1:
            coreset_anchor_distances = self._compute_batchwise_differences(features, features)
        else:
            
            # number of steps
            row_nb_step = int(np.ceil(len(features) / self.batch_size))

            # empty matrix
            coreset_anchor_distances = torch.empty(len(features), dtype=torch.float32)

            # Create a CUDA stream
            stream = torch.cuda.Stream()

            for r_step in tqdm.tqdm(np.arange(row_nb_step), desc='Computing batch row-wise differences...'):
                row_batch_start = r_step * self.batch_size
                row_batch_end = min((r_step + 1) * self.batch_size, len(features))
                
                with torch.cuda.stream(stream):
                    distance_matrix_b = torch.cdist(
                        features[row_batch_start: row_batch_end],
                        features,
                    ).norm(dim=1) 

                    # Synchronize the stream
                    torch.cuda.synchronize()

                coreset_anchor_distances[row_batch_start:row_batch_end] = distance_matrix_b.to('cpu', non_blocking=True)
    
        coreset_indices = []
        num_coreset_samples = int(len(features) * self.percentage)

        for _ in tqdm.tqdm(range(num_coreset_samples), desc='Selecting samples...'):
            select_idx = torch.argmax(coreset_anchor_distances).item()
            coreset_indices.append(select_idx)
            
            coreset_select_distance = torch.cdist(
                features,
                features[select_idx:select_idx+1],
            ).cpu()
            
            coreset_anchor_distances = torch.min(torch.cat([coreset_anchor_distances.unsqueeze(1), coreset_select_distance], dim=1), dim=1).values

        return np.array(coreset_indices)


class ApproximateGreedyCoresetSampler(GreedyCoresetSampler):
    def __init__(
        self,
        percentage: float,
        device: torch.device,
        seed: int = 0,
        number_of_starting_points: int = 10,
        dimension_to_project_features_to: int = 128,
    ):
        """Approximate Greedy Coreset sampling base class."""
        self.number_of_starting_points = number_of_starting_points
        super().__init__(
            percentage                       = percentage, 
            device                           = device, 
            seed                             = seed, 
            dimension_to_project_features_to = dimension_to_project_features_to
        )

    def _compute_greedy_coreset_indices(self, features: torch.Tensor) -> np.ndarray:
        """Runs approximate iterative greedy coreset selection.
        This greedy coreset implementation does not require computation of the
        full N x N distance matrix and thus requires a lot less memory, however
        at the cost of increased sampling times.
        Args:
            features: [NxD] input feature bank to sample.
        """
        number_of_starting_points = np.clip(
            self.number_of_starting_points, None, len(features)
        )
        start_points = np.random.choice(
            len(features), number_of_starting_points, replace=False
        ).tolist()

        approximate_distance_matrix = self._compute_batchwise_differences(
            features, features[start_points]
        )
        approximate_coreset_anchor_distances = torch.mean(
            approximate_distance_matrix, axis=-1
        ).reshape(-1, 1)
        coreset_indices = []
        num_coreset_samples = int(len(features) * self.percentage)

        for _ in tqdm.tqdm(range(num_coreset_samples), desc="Subsampling..."):
            select_idx = torch.argmax(approximate_coreset_anchor_distances).item()
            coreset_indices.append(select_idx)
            coreset_select_distance = self._compute_batchwise_differences(
                features, features[select_idx : select_idx + 1]  # noqa: E203
            )
            approximate_coreset_anchor_distances = torch.cat(
                [approximate_coreset_anchor_distances, coreset_select_distance],
                dim=-1,
            )
            approximate_coreset_anchor_distances = torch.min(
                approximate_coreset_anchor_distances, dim=1
            ).values.reshape(-1, 1)

        return np.array(coreset_indices)


class RandomSampler(BaseSampler):
    def __init__(self, percentage: float):
        super().__init__(percentage)

    def run(
        self, features: Union[torch.Tensor, np.ndarray]
    ) -> Union[torch.Tensor, np.ndarray]:
        """Randomly samples input feature collection.
        Args:
            features: [N x D]
        """
        num_random_samples = int(len(features) * self.percentage)
        subset_indices = np.random.choice(
            len(features), num_random_samples, replace=False
        )
        subset_indices = np.array(subset_indices)
        return features[subset_indices]