Feature Extraction - Refactor command-line script into a Python module

.gitignore

@@ -94,4 +94,7 @@ dmypy.json
 # misc
 *.mp4
 sweep*/
-core*
+core*
+
+features_outputs
+*.pth

translation/README.md

@@ -1,10 +1,11 @@
 ## Translation
 
 ### Feature Extraction
+
 Once you have obtained a pretrained model, and your data is prepared as outlined in [DATASETS.md](../DATASETS.md), you are ready for feature extraction. Below is an example for how it can be run.
 
- `num_items_per_shard` is tuned based on the size of our dataset (fewer items per shard result in larger slurm array).
- `max_batch_size=2` is based on our setup using fp32 and 32GB V100 GPUs. When using a CLIP checkpoint with SignHiera vision tower, pass `from_clip=true`. If you are not using a SLURM cluster, you can pass `launcher.cluster=local` to run feature extraction on the local machine.
+`num_items_per_shard` is tuned based on the size of our dataset (fewer items per shard result in larger slurm array).
+`max_batch_size=2` is based on our setup using fp32 and 32GB V100 GPUs. When using a CLIP checkpoint with SignHiera vision tower, pass `from_clip=true`. If you are not using a SLURM cluster, you can pass `launcher.cluster=local` to run feature extraction on the local machine.
 
 For additional parameters, please refer to the `FeatureExtractionConfig` dataclass in [launch_feature_extraction.py](launch_feature_extraction.py).
 
@@ -27,6 +28,7 @@ python launch_feature_extraction.py \
 ```
 
 After extracting the features, you should make sure they follow the nested folder structure expected by our [`SignFeaturesDataset`](../src/ssvp_slt/data/sign_features_dataset.py) below. The manifest `.tsv` files should have the following columns: `["video_name", "length", "label"]`, where `video_name` uniquely identifies a video in the dataset, `length` denotes the length of the features tensor, and `label` is the translation. You can obtain manifest files with the length column populated based on the extracted features via our provided [get_feature_lengths.py](../scripts/get_feature_lengths.py) script.
+
 ```
 ├── DATASET_NAME/
     ├── manifests/
@@ -38,16 +40,18 @@ After extracting the features, you should make sure they follow the nested folde
             ├── 00000 (prefix)
                 ├── 00000-abcdef.pt (video_name.pt)
                 └── ...
-            ├── 00001 
+            ├── 00001
                 ├── 00001-cdefgh.pt
-                └── ...  
+                └── ...
         └── ...  (if extracted more than 1 epoch)
 ```
 
 ### Training
+
 After you have successfully extracted the features, you can simply run the translation pipeline as follows.
 
 For T5:
+
 ```bash
 # sweep over 5 random seeds
 python run_translation.py \
@@ -81,12 +85,12 @@ python run_translation.py \
     common.fp16=true
 ```
 
-
 Adjust the number of GPUs and gradient accumulation steps based on your setup. These setups are based on fp32 training on 32GB V100 GPUs. You can also play with the max sequence length (`data.max_source_positions`) to reduce memory footprint. While we recommend `common.fp16=true` for BART training and the `slt_tarres_fairseq` configuration, `fp16` does not work with T5. We have not tried `bfloat16` training but encourage adding support for it.
 
 For all additional training and evaluation parameters, refer to the dataclasses in [run_translation.py](run_translation.py).
 
 ### Evaluation
+
 You can run the translation pipeline in eval mode to only perform evaluation on the validation and test sets.
 To do this, simply pass the appropriate overrides, e.g.:
 
@@ -100,4 +104,10 @@ python run_translation.py \
     common.load_model=/path/to/finetuned-t5/best_model.pth
 ```
 
+### Feature Extraction Demo
+
+To help you get started with the feature extraction process, we've provided a Colab notebook that demonstrates the key steps. The notebook walks through the workflow, allowing you to interactively extract features from your data using our tools.
+
+You can access the demo here: [Feature Extraction Demo on Colab](https://colab.research.google.com/drive/1EeL84ZrPRwaQBldPS9TqbWogJj52vn5Q?usp=sharing)
 
+Simply open the notebook, connect to a runtime, and follow the instructions. This example provides insights into how to utilize our feature extraction pipeline effectively.

translation/feature_extraction_module.py

@@ -0,0 +1,280 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from einops import rearrange
+from typing import Any, Generator, List, Union, Tuple
+
+from omegaconf import MISSING
+import pandas as pd
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from ssvp_slt.data.video_dataset import VideoDataset
+import ssvp_slt.modeling.sign_hiera as sign_hiera
+import ssvp_slt.util.misc as misc
+from ssvp_slt.modeling.sign_hiera import SignHiera
+from stopes.core import Requirements, StopesModule
+
+@dataclass
+class LauncherConfig:
+    cluster: str = "slurm"
+    """Cluster to use by stopes launcher, e.g. 'slurm' or 'local'"""
+
+    partition: str = "gpu"
+    """Partition to use in slurm environment"""
+
+    max_jobarray_jobs: int = 128
+    """Maximum number of jobs to launch in parallel"""
+
+@dataclass
+class FeatureExtractionConfig:
+    data_dir: str = MISSING
+    """Folder on disk containing manifests and videos"""
+
+    output_dir: str = MISSING
+    """Folder on disk where features will be saved"""
+
+    model_name: str = MISSING
+    """Hiera model identifier, e.g. hiera_base_128x224"""
+
+    pretrained_model_path: str = MISSING
+    """Path to trained feature extractor weights"""
+
+    from_clip: bool = False
+    """Whether to initialize weights from a trained CLIP model"""
+
+    split: str = "test"
+    """Dataset split, i.e. train, val or test"""
+
+    fp16: bool = False
+    """Whether to extract features in fp16 rather than fp32"""
+
+    do_aug: bool = False
+    """Whether to perform data augmentation"""
+
+    epochs: int = 1
+    """Number of epochs to extract (useful when performing augmentation)"""
+
+    num_frames: int = 128
+    """Number of video frames per clip"""
+
+    sampling_rate: int = 2
+    """Sample every n-th frame"""
+
+    target_fps: int = 25
+    """Target framerate in frames per second"""
+
+    num_items_per_shard: int = 50
+    """Number of videos that will be processed sequentially in each shard"""
+
+    max_batch_size: int = 2
+    """Maximum batch size (in 128-frame clips). Videos longer than this will be processed in shards."""
+
+    video_backend: str = "pyav"
+    """Video backend. Should match the video backend used during pretraining."""
+
+    launcher: LauncherConfig = LauncherConfig()
+    """Stopes launcher configuration"""
+
+def shard_generator(data: Any, shard_size: int) -> Generator[Any, None, None]:
+    for i in range(0, len(data), shard_size):
+        yield data[i : i + shard_size]
+
+class FeatureExtractionModule(StopesModule):
+    def __init__(self, config: FeatureExtractionConfig):
+        super().__init__(config, FeatureExtractionConfig)
+
+        manifest_file = Path(self.config.data_dir) / "manifests" / f"{self.config.split}.tsv"
+
+        self.num_items = len(
+            pd.read_csv(
+                manifest_file,
+                delimiter="\t",
+                names=["video_name", "duration", "caption"],
+                quoting=3,
+            )
+        )
+
+        Path(self.config.output_dir).mkdir(exist_ok=True, parents=True)
+
+    def requirements(self) -> Requirements:
+        """Requirements for each submitted job"""
+        return Requirements(
+            nodes=1,
+            mem_gb=6,
+            tasks_per_node=1,
+            gpus_per_node=1,
+            cpus_per_task=1,
+            timeout_min=60 * 72,
+        )
+
+    def name(self) -> str:
+        name = (
+            f"feature_extractor_{self.config.split}_{self.config.model_name}_"
+            f"{Path(self.config.pretrained_model_path).stem}"
+        )
+        return name
+
+    @property
+    def num_shards(self) -> int:
+        return math.ceil(self.num_items / self.config.num_items_per_shard)
+
+    def array(self) -> List[Any]:
+        return [1] * self.num_shards
+
+    def load_model(self, device: Union[torch.device, str] = "cuda") -> SignHiera:
+        """
+        Loads a pretrained SignHiera model for feature extraction and moves it to specified device
+        """
+
+        if self.config.from_clip:
+            if "hiera" in self.config.model_name:
+                model = SignHiera.from_clip_model(
+                    self.config.model_name, self.config.pretrained_model_path
+                )
+            else:
+                raise ValueError(
+                    f"Loading `{self.config.model_name}` from a CLIP model is not supported."
+                )
+        else:
+            model = sign_hiera.__dict__[self.config.model_name](pretrained=True, strict=False)
+            misc.load_model(model, self.config.pretrained_model_path)
+
+        model.head = nn.Identity()
+
+        print(f"Number of parameters: {sum([p.numel() for p in model.parameters()])}")
+
+        model.eval()
+        model.to(device)
+
+        return model
+
+    def get_dataloader(self, indices: Tuple[int, int]) -> DataLoader:
+        """
+        Creates a dataloader for video feature extraction
+        Uses a fixed T//2 stride for the sliding window
+        """
+
+        if self.config.video_backend == "cuda":
+            torch.multiprocessing.set_start_method("spawn")
+
+        dataset = VideoDataset(
+            mode=self.config.split,
+            video_backend=self.config.video_backend,
+            target_fps=self.config.target_fps,
+            data_dir=self.config.data_dir,
+            sampling_rate=self.config.sampling_rate,
+            num_frames=self.config.num_frames,
+            rand_aug=self.config.do_aug,
+            train_random_horizontal_flip=self.config.do_aug,
+            train_random_crop=self.config.do_aug,
+            feature_extraction=True,
+            feature_extraction_stride=self.config.num_frames // 2,
+            indices=indices,
+            gpu=(torch.cuda.current_device() if self.config.video_backend == "cuda" else None),
+        )
+
+        dataloader = DataLoader(
+            dataset,
+            batch_size=1,
+            num_workers=1 if self.config.video_backend == "cuda" else 2,
+            persistent_workers=True,
+            pin_memory=not self.config.video_backend == "cuda",
+        )
+
+        return dataloader
+
+    def run(self, iteration_value: Any, iteration_index: int) -> None:
+        """
+        Extract features for a shard of the data
+        """
+        output_dir = Path(self.config.output_dir)
+
+        start_id, end_id = (
+            self.config.num_items_per_shard * iteration_index,
+            self.config.num_items_per_shard * (iteration_index + 1),
+        )
+        print(f"Indices: {start_id}-{end_id}")
+
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"Using device {device}")
+
+        model = self.load_model(device=device)
+        dataloader = self.get_dataloader(indices=(start_id, end_id))
+
+        # Split data into folders using prefix based on their video name
+        # Helps reduce load when file system is NFS
+        fids = [
+            Path(dataloader.dataset.path_to_videos[i]).stem for i in range(len(dataloader.dataset))
+        ]
+        prefixes = [fid[:5] for fid in fids]
+        for epoch in tqdm(range(self.config.epochs), desc="Creating folder structure"):
+            for prefix in prefixes:
+                prefix_path = output_dir / str(epoch) / prefix
+                prefix_path.mkdir(parents=True, exist_ok=True)
+
+        for epoch in range(self.config.epochs):
+            print(f"Epoch: {epoch}")
+
+            prefetcher = misc.Prefetcher(dataloader, device=device)
+
+            start_time = time.time()
+            idx = 0
+            pbar = tqdm(total=len(dataloader))
+
+            batch = next(prefetcher)
+            while batch is not None:
+                frames = batch["frames"].float()
+                padding = batch["padding"]
+
+                if frames.dim() == 6:
+                    frames = rearrange(frames, "b r c t h w -> (b r) c t h w")
+                if padding.dim() == 2:
+                    padding = rearrange(padding, "b r -> (b r)")
+
+                # Process in chunks
+                if len(frames) > self.config.max_batch_size:
+                    shard_outputs = []
+                    frames_shards = shard_generator(frames, self.config.max_batch_size)
+                    padding_shards = shard_generator(padding, self.config.max_batch_size)
+
+                    for frames_shard, padding_shard in zip(frames_shards, padding_shards):
+
+                        with torch.inference_mode(), torch.cuda.amp.autocast(
+                            enabled=self.config.fp16
+                        ):
+                            shard_output = model.extract_features(
+                                frames_shard, padding=padding_shard
+                            ).cpu()
+                        if len(shard_output.shape) == 1:
+                            shard_output = shard_output.unsqueeze(0)
+                        shard_outputs.append(shard_output)
+
+                    outputs = torch.concatenate(shard_outputs, dim=0)
+
+                else:
+                    with torch.inference_mode(), torch.cuda.amp.autocast(enabled=self.config.fp16):
+                        outputs = model.extract_features(frames, padding=padding).detach().cpu()
+
+                fid = fids[idx]
+                prefix = prefixes[idx]
+
+                # Save to disk
+                output_file = output_dir / str(epoch) / prefix / f"{fid}.pt"
+                torch.save(outputs, output_file)
+
+                idx += 1
+                pbar.update(1)
+                batch = next(prefetcher)
+
+            print(f"Epoch time: {time.time() - start_time:.2f}s")
+            pbar.close()

translation/feature_extraction_test.py

@@ -0,0 +1,39 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+from feature_extraction_module import FeatureExtractionConfig, FeatureExtractionModule, LauncherConfig
+from unittest.mock import patch, MagicMock
+from utils.download_model import download_model
+
+class TestFeatureExtractionModule(unittest.TestCase):
+    def setUp(self):
+        # Mock the configuration for the FeatureExtractionModule
+        model_path = download_model('https://dl.fbaipublicfiles.com/SONAR/asl/signhiera_mock.pth')
+        self.config = FeatureExtractionConfig(
+            #You need to download blurred dailymoth-70h dataset first
+            data_dir="path/to/your/blurred/dataset",
+            pretrained_model_path=model_path, 
+            launcher=LauncherConfig(cluster="local")
+        )
+        self.module = FeatureExtractionModule(self.config)
+
+    @patch("torch.cuda.is_available", return_value=False)  # Mock CUDA for CPU
+    def test_load_model(self, mock_cuda):
+        # Test if the model loads properly
+        model = self.module.load_model()
+        self.assertIsNotNone(model)
+
+    @patch("ssvp_slt.data.video_dataset.VideoDataset")
+    def test_get_dataloader(self, mock_video_dataset):
+        # Mock the VideoDataset and test the dataloader
+        mock_dataset = MagicMock()
+        mock_video_dataset.return_value = mock_dataset
+        dataloader = self.module.get_dataloader((0, 10))
+        self.assertIsNotNone(dataloader)
+
+if __name__ == "__main__":
+    unittest.main()