Benchmarking script for image pipeline

benchmarking/nightly-benchmark.yaml

@@ -33,7 +33,7 @@ datasets:
   - name: "mscoco"
     formats:
     - type: "wds"
-      path: "{datasets_path}/mscoco/wds/truncated_100K_mscoco_benchmarking"
+      path: "{datasets_path}/mscoco/wds/full_mscoco_benchmarking"
   - name: "mscoco_model_weights"
     formats:
     - type: "files"
@@ -437,15 +437,15 @@ entries:
 
   - name: image_curation
     enabled: true
-    script: "{curator_repo_dir}/tutorials/image/getting-started/image_curation_example.py"
+    script: image_pipeline_benchmark.py
     args: >-
       --input-wds-dataset-dir {dataset:mscoco,wds}
       --output-dataset-dir {session_entry_dir}/scratch/output
       --model-dir {dataset:mscoco_model_weights,files}
-      --batch-size 100
-      --embedding-batch-size 100
-      --aesthetic-batch-size 100
-      --nsfw-batch-size 100
+      --batch-size 1000
+      --embedding-batch-size 500
+      --aesthetic-batch-size 500
+      --nsfw-batch-size 500
       --tar-files-per-partition 10
       --aesthetic-threshold 0.9
       --nsfw-threshold 0.9

benchmarking/scripts/image_pipeline_benchmark.py

@@ -0,0 +1,277 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  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.
+
+import argparse
+import os
+import time
+
+from nemo_curator.backends.xenna import XennaExecutor
+from nemo_curator.core.client import RayClient
+from nemo_curator.pipeline import Pipeline
+from nemo_curator.stages.file_partitioning import FilePartitioningStage
+from nemo_curator.stages.image.embedders.clip_embedder import ImageEmbeddingStage
+from nemo_curator.stages.image.filters.aesthetic_filter import ImageAestheticFilterStage
+from nemo_curator.stages.image.filters.nsfw_filter import ImageNSFWFilterStage
+from nemo_curator.stages.image.io.image_reader import ImageReaderStage
+from nemo_curator.stages.image.io.image_writer import ImageWriterStage
+
+
+def create_image_curation_pipeline(args: argparse.Namespace) -> Pipeline:
+    """Create image curation pipeline with file partitioning, image reading, embedding, aesthetic scoring, and NSFW detection stages."""
+
+    # Define pipeline
+    pipeline = Pipeline(name="image_curation", description="Curate images with embeddings and quality scoring")
+
+    # Stage 0: Partition tar files for parallel processing
+    pipeline.add_stage(FilePartitioningStage(
+        file_paths=args.input_wds_dataset_dir,
+        files_per_partition=args.tar_files_per_partition,
+        file_extensions=[".tar"],
+    ))
+
+    # Stage 1: Read images from webdataset tar files (now runs in parallel)
+    pipeline.add_stage(ImageReaderStage(
+        batch_size=args.batch_size,
+        verbose=args.verbose,  # Force verbose to see debug info
+        num_threads=16,  # More threads for I/O
+        num_gpus_per_worker=0.25,
+    ))
+
+    # Stage 2: Generate CLIP embeddings for images
+    pipeline.add_stage(ImageEmbeddingStage(
+        model_dir=args.model_dir,
+        num_gpus_per_worker=args.embedding_gpus_per_worker,
+        model_inference_batch_size=args.embedding_batch_size,
+        remove_image_data=False,
+        verbose=args.verbose,
+    ))
+
+    # Stage 3: Generate aesthetic quality scores and filter
+    pipeline.add_stage(ImageAestheticFilterStage(
+        model_dir=args.model_dir,
+        num_gpus_per_worker=args.aesthetic_gpus_per_worker,
+        model_inference_batch_size=args.aesthetic_batch_size,
+        score_threshold=args.aesthetic_threshold,
+        verbose=args.verbose,
+    ))
+
+    # Stage 4: Generate NSFW probability scores and filter
+    pipeline.add_stage(ImageNSFWFilterStage(
+        model_dir=args.model_dir,
+        num_gpus_per_worker=args.nsfw_gpus_per_worker,
+        model_inference_batch_size=args.nsfw_batch_size,
+        score_threshold=args.nsfw_threshold,
+        verbose=args.verbose,
+    ))
+
+    # Stage 5: Write down to disk
+    pipeline.add_stage(ImageWriterStage(
+        output_dir=args.output_dataset_dir,
+        images_per_tar=args.images_per_tar,
+        remove_image_data=True,
+        verbose=args.verbose,
+    ))
+
+    return pipeline
+
+
+def main(args: argparse.Namespace) -> None:
+    """Main execution function for image curation pipeline."""
+
+    ray_client = RayClient()
+    ray_client.start()
+
+    print("Starting image curation pipeline...")
+    print(f"Input parquet file: {args.input_parquet}")
+    print(f"Input webdataset directory: {args.input_wds_dataset_dir}")
+    print(f"Output webdataset directory: {args.output_dataset_dir}")
+    print(f"Model directory: {args.model_dir}")
+    print(f"Tar files per partition: {args.tar_files_per_partition}")
+    print(f"Task batch size: {args.batch_size}")
+    print("\n" + "=" * 50 + "\n")
+
+    # Step 1: Download and prepare webdataset from parquet file
+    if not args.skip_download:
+        assert False, "Downloading is not supported in pipeline benchmark."
+    else:
+        print("Step 1: Skipping download (using existing dataset)")
+        print(f"Using existing dataset at: {args.input_wds_dataset_dir}")
+        print("\n" + "=" * 50 + "\n")
+
+    # Step 2: Create and run curation pipeline
+    print("Step 2: Running image curation pipeline...")
+    start_time = time.time()
+    pipeline = create_image_curation_pipeline(args)
+
+    # Print pipeline description
+    print(pipeline.describe())
+    print("\n" + "=" * 50 + "\n")
+
+    # Create executor
+    executor = XennaExecutor()
+
+    # Execute pipeline
+    pipeline.run(executor)
+
+    end_time = time.time()
+
+    # Calculate and print execution time
+    execution_time = end_time - start_time
+    hours, remainder = divmod(execution_time, 3600)
+    minutes, seconds = divmod(remainder, 60)
+
+    print("\nImage curation pipeline completed!")
+    print(f"Total execution time: {int(hours):02d}:{int(minutes):02d}:{seconds:.2f}")
+    print(f"Total execution time: {execution_time:.2f} seconds")
+    print(f"\nProcessed dataset available at: {args.output_dataset_dir}")
+
+    ray_client.stop()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Image curation pipeline with embedding generation and quality scoring"
+    )
+
+    # Dataset arguments
+    parser.add_argument(
+        "--input-parquet",
+        type=str,
+        required=False,
+        default=None,
+        help="Path to input parquet file containing image URLs and metadata"
+    )
+    parser.add_argument(
+        "--input-wds-dataset-dir",
+        type=str,
+        required=True,
+        help="Directory to save the downloaded webdataset"
+    )
+    parser.add_argument(
+        "--output-dataset-dir",
+        type=str,
+        required=True,
+        help="Directory to save the resulting webdataset"
+    )
+    parser.add_argument(
+        "--download-processes",
+        type=int,
+        default=8,
+        help="Number of parallel processes for downloading images"
+    )
+    parser.add_argument(
+        "--entries-per-tar",
+        type=int,
+        default=1000,
+        help="Number of entries per tar shard during download"
+    )
+    parser.add_argument(
+        "--skip-download",
+        action="store_true",
+        default=False,
+        help="Skip dataset download and use existing webdataset"
+    )
+
+    # Image reader arguments
+    parser.add_argument(
+        "--tar-files-per-partition",
+        type=int,
+        default=1,
+        help="Number of tar files to process per partition (controls parallelism) for FilePartitioningStage"
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=100,
+        help="Number of images per ImageBatch for the reader stage"
+    )
+
+    # General arguments
+    parser.add_argument(
+        "--model-dir",
+        type=str,
+        required=True,
+        help="Path to model directory containing all model weights"
+    )
+    parser.add_argument(
+        "--verbose",
+        action="store_true",
+        default=False,
+        help="Enable verbose logging for all stages"
+    )
+
+    # Embedding stage arguments
+    parser.add_argument(
+        "--embedding-batch-size",
+        type=int,
+        default=32,
+        help="Batch size for embedding generation"
+    )
+    parser.add_argument(
+        "--embedding-gpus-per-worker",
+        type=float,
+        default=0.25,
+        help="GPU allocation per worker for embedding generation"
+    )
+
+    # Aesthetic scoring arguments
+    parser.add_argument(
+        "--aesthetic-batch-size",
+        type=int,
+        default=32,
+        help="Batch size for aesthetic scoring"
+    )
+    parser.add_argument(
+        "--aesthetic-gpus-per-worker",
+        type=float,
+        default=0.25,
+        help="GPU allocation per worker for aesthetic scoring"
+    )
+    parser.add_argument(
+        "--aesthetic-threshold",
+        type=float,
+        default=0.5,
+        help="Aesthetic score threshold for filtering (images below this score will be filtered out)"
+    )
+
+    # NSFW scoring arguments
+    parser.add_argument(
+        "--nsfw-batch-size",
+        type=int,
+        default=32,
+        help="Batch size for NSFW scoring"
+    )
+    parser.add_argument(
+        "--nsfw-gpus-per-worker",
+        type=float,
+        default=0.25,
+        help="GPU allocation per worker for NSFW scoring"
+    )
+    parser.add_argument(
+        "--nsfw-threshold",
+        type=float,
+        default=0.5,
+        help="NSFW score threshold for filtering (images above this score will be filtered out as NSFW)"
+    )
+
+    # Output dataset arguments
+    parser.add_argument(
+        "--images-per-tar",
+        type=int,
+        default=100,
+        help="Number of images per tar file in output dataset"
+    )
+
+    args = parser.parse_args()
+    main(args)