DOCSEARCHbyAlgolia
Submit Docs

SemDeDup

repository·main·Indexed 41 days ago

https://github.com/facebookresearch/semdedup

A method for identifying and removing semantically similar data pairs from large-scale datasets to reduce redundancy and speed up the learning process. The pipeline includes computing embeddings via get_embeddings(), performing K-Means clustering with compute_centroids(), sorting clusters, and extracting pruned data.

Tokens
2.1K
Snippets
7
Records
7
Agent score
50%

What's inside SemDeDup

  1. Run K-Means clustering on embeddings

    main

    Perform K-Means clustering on the embeddings generated in step 1. You must first modify the configuration file located at configs/openclip/clustering_configs.yaml.

    You can run this via the CLI using submitit for job submission, or directly via the compute_centroids Python function.

    # Via CLI
python compute_centroids.py --confg-file "/configs/openclip/clustering_configs.yaml" \
                            --partition <partition> \
                            --ngpus 1 \
                            --cpus-per-task 10 \
                            --timeout 300

  2. Install SemDeDup via Conda

    main

    To set up the SemDeDup environment, use the provided semdedup_conda_env.yml file to create a new conda environment. This ensures all dependencies are correctly configured.

    conda env create -n semdedup --file semdedup_conda_env.yml 
conda activate semdedup

  3. Compute and store embeddings with get_embeddings()

    main

    The first step in the pipeline is to compute embeddings from a pretrained model and store them in memory-mapped files.

    Your dataloader must return a tuple containing:

    1. data_batch: A batch of data examples.
    2. paths_batch: A unique identifier for the example (e.g., a file path like "n04235860_14959.JPEG").
    3. batch_indices: The global index for each example (integer between 0 and dataset_size - 1).

    You must provide two numpy memmap objects: emd_memmap (to store embeddings) and paths_memmap (to store paths).

    from compute_pretrained_embeddings import get_embeddings
import numpy as np

model = ...
dataloader = ...

path_str_type = ...
emb_memory_loc = ...
paths_memory_loc = ...
dataset_size = ...
emb_size = ...

# Initialize memmaps
emb_array = np.memmap(emb_memory_loc, dtype='float32', mode='w+', shape=(dataset_size, emb_size))
path_array = np.memmap(paths_memory_loc, dtype=path_str_type, mode='w+', shape=(dataset_size,))

get_embeddings(model, dataloader, emb_array, path_array)

  4. Sort clusters with assign_and_sort_clusters()

    main

    After clustering, use assign_and_sort_clusters to organize the clusters. This can be done via the sort_clusters.py CLI script or the Python API.

    from sort_clusters import assign_and_sort_clusters

# Assumes emb_memory, paths_memory, and params are already defined
assign_and_sort_clusters(
    data=emb_memory,
    paths_list=paths_memory,
    sim_metric=params["sim_metric"],
    keep_hard=params["keep_hard"],
    kmeans_with_cos_dist=params["kmeans_with_cos_dist"],
    save_folder=params["save_folder"],
    sorted_clusters_file_loc=params["sorted_clusters_file_loc"],
    cluster_ids=range(0, params["ncentroids"]),
    logger=logger,
)

  5. Use compute_centroids() for K-Means clustering

    main

    The compute_centroids function from the clustering.clustering module performs K-Means clustering. It requires embeddings loaded as a numpy memmap and various parameters from a configuration file.

    from clustering.clustering import compute_centroids
import numpy as np
import yaml
import logging

# Setup logger and load config
logger = logging.getLogger(__name__)
logger.addHandler(logging.StreamHandler())
confg_file = "configs/openclip/clustering_configs.yaml"

with open(confg_file, 'r') as y_file:
    params = yaml.load(y_file, Loader=yaml.FullLoader)

# Load memmaps
emb_memory = np.memmap(params['emb_memory_loc'], dtype='float32', mode='r', shape=(params['dataset_size'], params['emb_size']))
paths_memory = np.memmap(params['paths_memory_loc'], dtype=params['path_str_type'], mode='r', shape=(params['dataset_size'],))

compute_centroids(
    data=emb_memory,
    ncentroids=params['ncentroids'],
    niter=params['niter'],
    seed=params['seed'],
    Kmeans_with_cos_dist=params['Kmeans_with_cos_dist'],
    save_folder=params['save_folder'],
    logger=logger,
    verbose=True,
)

  6. Extract deduplicated data with extract_pruned_data()

    main

    The final step is to extract the paths of the examples that were kept after the SemDeDup pruning process. This produces a text file containing the paths to the deduplicated data.

    from extract_dedup_data import extract_pruned_data

# Parameters required:
# sorted_clusters_path: Output from step 3
# semdedup_pruning_tables_path: Output from step 4
# eps: The epsilon value used in step 4
# num_clusters: Number of clusters from K-Means
# output_txt_path: Destination for the resulting text file

extract_pruned_data(
    sorted_clusters_path, 
    semdedup_pruning_tables_path, 
    eps, 
    num_clusters, 
    output_txt_path, 
    retreive_kept_samples=True
)

  7. Run the SemDeDup pipeline

    main

    The core SemDeDup process is executed via submit_semdedup_job.py. This script uses submitit to launch parallel jobs across multiple nodes. You must configure semdedup_configs.yaml before running.

    Key arguments:

    • --config-file: Path to the configuration file.
    • --eps-list: A list of epsilon values. Lower epsilon values result in more aggressive pruning.
    • --partition: Partition identifier.
    • --nodes: Number of nodes to use.
    • --tasks-per-node: Number of tasks per node to increase parallelization.
    • --cpus-per-task: Number of CPUs allocated per task.
    python submit_semdedup_job.py  --config-file "semdedup_configs.yaml" \
                               --eps-list <eps-list> \
                               --partition <partition> \
                               --nodes 8 \
                               --tasks-per-node 20 \
                               --cpus-per-task 4