Make faiss an optional dependency of VISSL and switch to faiss-gpu (f…

…acebookresearch#203) Summary: Pull Request resolved: facebookresearch#203 - making faiss an optional dependency. faiss is used only in 1 place in VISSL for clusterfit. - When user uses clusterfit, we check if faiss is installed or not. we give instructions to user at runtime to how to install faiss if not already done - we shift to using faiss-gpu since that's official dependency to use for faiss Reviewed By: bottler Differential Revision: D26725011 fbshipit-source-id: eac86595da9c998230aa74a4a00c7a894b6200f7
ArrowLuo · Mar 1, 2021 · 44bb572 · 44bb572
1 parent 701172a
commit 44bb572
Show file tree

Hide file tree

Showing 7 changed files with 29 additions and 5 deletions.
diff --git a/dev/packaging/vissl_conda/vissl/meta.yaml b/dev/packaging/vissl_conda/vissl/meta.yaml
@@ -21,7 +21,7 @@ requirements:
     - scikit-learn
     - parameterized
     - numpy >=1.11
-    - faiss
+    - faiss-gpu
     - fvcore
     - iopath
     - importlib_resources

diff --git a/docs/requirements.txt b/docs/requirements.txt
@@ -11,7 +11,7 @@ mock
 opencv-python
 tensorboard==1.15.0
 hydra-core>=1.0
-faiss>=1.5.3
+faiss-gpu
 cython
 scikit-learn
 parameterized==0.7.4

diff --git a/requirements.txt b/requirements.txt
@@ -1,6 +1,5 @@
 tensorboard==1.15.0
 hydra-core>=1.0
-faiss>=1.5.3
 cython
 scikit-learn
 parameterized==0.7.4

diff --git a/setup.py b/setup.py
@@ -55,6 +55,7 @@ def get_version():
             "pre-commit",
             "nbconvert",
             "bs4",
+            "faiss-gpu",
         ]
     },
 )
diff --git a/tools/cluster_features_and_label.py b/tools/cluster_features_and_label.py
@@ -10,7 +10,6 @@
 from argparse import Namespace
 from typing import Any, List
 
-import faiss
 import numpy as np
 from hydra.experimental import compose, initialize_config_module
 from vissl.data import build_dataset
@@ -21,7 +20,7 @@
 from vissl.utils.hydra_config import AttrDict, convert_to_attrdict, is_hydra_available
 from vissl.utils.io import save_file
 from vissl.utils.logger import setup_logging, shutdown_logging
-from vissl.utils.misc import merge_features, set_seeds
+from vissl.utils.misc import merge_features, set_seeds, is_faiss_available
 
 
 def get_data_features_and_images(cfg: AttrDict):
@@ -45,6 +44,13 @@ def get_data_features_and_images(cfg: AttrDict):
 
 
 def cluster_features_and_label(args: Namespace, cfg: AttrDict):
+    # faiss is an optional dependency for VISSL.
+    assert is_faiss_available(), (
+        "Please install faiss using conda install faiss-gpu -c pytorch "
+        "if using conda or pip install faiss-gpu"
+    )
+    import faiss
+
     cluster_backend = cfg.CLUSTERFIT.CLUSTER_BACKEND
     num_clusters = cfg.CLUSTERFIT.NUM_CLUSTERS
     data_split = cfg.CLUSTERFIT.FEATURES.DATA_PARTITION

diff --git a/vissl/data/datasets/coco.py b/vissl/data/datasets/coco.py
@@ -53,6 +53,7 @@ def get_valid_objs(entry, objs):
 
 
 def get_coco_imgs_labels_info(split, data_source_dir, args):
+    # pycocotools is an optional dependency for VISSL
     from pycocotools.coco import COCO
 
     json_file = f"{data_source_dir}/annotations/instances_{split}2014.json"

diff --git a/vissl/utils/misc.py b/vissl/utils/misc.py
@@ -26,6 +26,23 @@ def is_fairscale_sharded_available():
     return fairscale_sharded_available
 
 
+def is_faiss_available():
+    """
+    Check if faiss is available with simple python imports.
+
+    To install faiss, simply do:
+        If using PIP env: `pip install faiss-gpu`
+        If using conda env: `conda install faiss-gpu -c pytorch`
+    """
+    try:
+        import faiss  # NOQA
+
+        faiss_available = True
+    except ImportError:
+        faiss_available = False
+    return faiss_available
+
+
 def is_opencv_available():
     """
     Check if opencv is available with simple python imports.