class_docstr: Utility class to deal with features extracted with extract_engine
  
  For multi-gpu feature extraction, each GPU saves features corresponding to its
  share of the data. This class offers an API to abstract away the loading of
  these extracted features.
class_docstr_tok: ["Utility","class","to","deal","with","features","extracted","with","extract_engine","For","multi-gpu","feature","extraction",",","each","GPU","saves","features","corresponding","to","its","share","of","the","data",".","This","class","offers","an","API","to","abstract","away","the","loading","of","these","extracted","features","."]
docstr: This function sample N features across all GPUs in an optimized way, using
  reservoir sampling, to avoid loading all features in memory.
  
  This is especially useful if the number of feature is huge, cannot hold into
  memory, and we need a small number of them to do an estimation (example of
  k-means on a 1B dataset: we can use a few random million samples to compute
  relatively good centroids)
  
  Args:
      input_dir (str): input path where the features are dumped
      split (str): whether the features are train or test data features
      layer (str): the features correspond to what layer of the model
      num_samples (int): how many features to sample, if negative load everything
      seed (int): the random seed used for sampling
      flatten_features (bool): whether or not to flatten the features
  
  Returns:
      output (Dict): contains features, targets, inds as the keys
docstr_tok: ["This","function","sample","N","features","across","all","GPUs","in","an","optimized","way",",","using","reservoir","sampling",",","to","avoid","loading","all","features","in","memory",".","This","is","especially","useful","if","the","number","of","feature","is","huge",",","can","not","hold","into","memory",",","and","we","need","a","small","number","of","them","to","do","an","estimation","(","example","of","k-means","on","a","1B","dataset",":","we","can","use","a","few","random","million","samples","to","compute","relatively","good","centroids",")","Args",":","input_dir","(","str",")",":","input","path","where","the","features","are","dumped","split","(","str",")",":","whether","the","features","are","train","or","test","data","features","layer","(","str",")",":","the","features","correspond","to","what","layer","of","the","model","num_samples","(","int",")",":","how","many","features","to","sample",",","if","negative","load","everything","seed","(","int",")",":","the","random","seed","used","for","sampling","flatten_features","(","bool",")",":","whether","or","not","to","flatten","the","features","Returns",":","output","(","Dict",")",":","contains","features",",","targets",",","inds","as","the","keys"]
code: def sample_features(
      cls,
      input_dir: str,
      split: str,
      layer: str,
      num_samples: int,
      seed: int,
      flatten_features: bool = False,
  ):
      """
      This function sample N features across all GPUs in an optimized way, using
      reservoir sampling, to avoid loading all features in memory.
  
      This is especially useful if the number of feature is huge, cannot hold into
      memory, and we need a small number of them to do an estimation (example of
      k-means on a 1B dataset: we can use a few random million samples to compute
      relatively good centroids)
  
      Args:
          input_dir (str): input path where the features are dumped
          split (str): whether the features are train or test data features
          layer (str): the features correspond to what layer of the model
          num_samples (int): how many features to sample, if negative load everything
          seed (int): the random seed used for sampling
          flatten_features (bool): whether or not to flatten the features
  
      Returns:
          output (Dict): contains features, targets, inds as the keys
      """
  
      if num_samples < 0:
          return cls.load_features(
              input_dir=input_dir,
              split=split,
              layer=layer,
              flatten_features=flatten_features,
          )
  
      features = []
      targets = []
      indices = []
  
      # Find the shards containing the features to samples
      count = 0
      shard_paths = cls.get_shard_file_names(
          input_dir, split=split, layer=layer
      )
      if not shard_paths:
          raise ValueError(f"No features found for {split} {layer}")
  
      # Use reservoir sampling to sample some features
      rng = np.random.default_rng(seed)
      for shard_path in shard_paths:
          shard_content = cls.load_feature_shard(shard_path)
          for idx in range(shard_content.num_samples):
              count += 1
  
              # Fill the reservoir of samples
              if len(features) < num_samples:
                  features.append(shard_content.features[idx])
                  targets.append(shard_content.targets[idx])
                  indices.append(shard_content.indices[idx])
  
              # Else implement reservoir sampling substitution
              else:
                  pos = rng.integers(low=0, high=count)
                  if pos < num_samples:
                      features[pos] = shard_content.features[idx]
                      targets[pos] = shard_content.targets[idx]
                      indices[pos] = shard_content.indices[idx]
  
      # Cast the output to numpy arrays
      features = np.stack(features)
      targets = np.stack(targets)
      indices = np.stack(indices)
  
      # Sort the entries by sample index
      sorted_indices = np.argsort(indices)
      indices = indices[sorted_indices]
      features = features[sorted_indices]
      targets = targets[sorted_indices]
  
      # Return the output
      if flatten_features:
          N = len(indices)
          features = features.reshape(N, -1)
      output = {
          "features": features,
          "targets": targets,
          "inds": indices,
      }
      logging.info(f"Features: {output['features'].shape}")
      logging.info(f"Targets: {output['targets'].shape}")
      logging.info(f"Indices: {output['inds'].shape}")
      return output
  
file_path: vissl/vissl/utils/extract_features_utils.py
imports: import logging
  import os
  import re
  from typing import List, NamedTuple
  import numpy
  from iopath.common.file_io import g_pathmgr
  from vissl.utils.io import load_file, makedir
label_desc: Use image node_id 5 for calling the ExtractedFeaturesLoader obj's underlying member method code with example usage: obj.sample_features(cls, input_dir, split, layer, num_samples, seed, flatten_features) and returns: output, cls