Repository: Lightning-Universe/lightning-bolts

Issue:
NameError: name 'shuffle' is not defined
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

### To Reproduce

Use `ImagenetDataModule` without having `scikit-learn` installed.

Inside `imagenet_dataset.py` there is the following try-except in the imports:
```
try:
    from sklearn.utils import shuffle
except ImportError:
    warn('You want to use `sklearn` which is not installed yet,'  # pragma: no-cover
         ' install it with `pip install sklearn`.')
```
so the program flow continues although sklearn might not be available.

However, the `shuffle` is used in the `UnlabeledImagenet` constructor. So in case sklearn is actually missing, it will crash with the following trace:

```
  File "/home/andrea.palazzi/.local/lib/python3.7/site-packages/pytorch_lightning/accelerators/base_accelerator.py", line 50, in train_or_test
    results = self.trainer.train()
  File "/home/andrea.palazzi/.local/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 476, in train
    self.run_sanity_check(self.get_model())
  File "/home/andrea.palazzi/.local/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 656, in run_sanity_check
    self.reset_val_dataloader(ref_model)
  File "/home/andrea.palazzi/.local/lib/python3.7/site-packages/pytorch_lightning/trainer/data_loading.py", line 316, in reset_val_dataloader
    self.num_val_batches, self.val_dataloaders = self._reset_eval_dataloader(model, 'val')
  File "/home/andrea.palazzi/.local/lib/python3.7/site-packages/pytorch_lightning/trainer/data_loading.py", line 236, in _reset_eval_dataloader
    dataloaders = self.request_dataloader(getattr(model, loader_name))
  File "/home/andrea.palazzi/.local/lib/python3.7/site-packages/pytorch_lightning/trainer/data_loading.py", line 339, in request_dataloader
    dataloader = dataloader_fx()
  File "/home/andrea.palazzi/.local/lib/python3.7/site-packages/pl_bolts/datamodules/imagenet_datamodule.py", line 159, in val_dataloader
    transform=transforms)
  File "/home/andrea.palazzi/.local/lib/python3.7/site-packages/pl_bolts/datamodules/imagenet_dataset.py", line 74, in __init__
    self.imgs = shuffle(self.imgs, random_state=1234)
NameError: name 'shuffle' is not defined

```

### Expected behavior

I would expect `scikit-learn` to be installed automatically when installing bolts, as it's not really optional. 

Or at least, the program could fail at import time, it would still be easier to understand where the error comes from. The current way it's not obvious, because the warning gets mixed up among other info/warnings I might have on console and I found myself wondering why `shuffle` is not defined.

### Environment
- torch-related: 
  - pytorch-lightning==0.10.0
  - pytorch-lightning-bolts==0.2.2
  - torch==1.6.0+cu101
  - torchvision==0.7.0+cu101
 - OS (e.g., Linux): Ubuntu 20.04.1 LTS
 - How you installed PyTorch (`conda`, `pip`, source): pip
 - Python version: 3.7.5
 - CUDA/cuDNN version: 101

Thanks.



Discussion:
Hi! thanks for your contribution!, great first issue!
@ndrplz Thank you for reporting the issue in detail! I agree that there should be something like `ModuleNotFoundError` to inform users of required modules e.g. `sklearn`.

@Borda What do you think? Can I submit a PR to fix this?

Resolving patch:
diff --git a/pl_bolts/datasets/imagenet_dataset.py b/pl_bolts/datasets/imagenet_dataset.py
--- a/pl_bolts/datasets/imagenet_dataset.py
+++ b/pl_bolts/datasets/imagenet_dataset.py
@@ -1,5 +1,6 @@
 import gzip
 import hashlib
+import importlib
 import os
 import shutil
 import tarfile
@@ -11,9 +12,10 @@
 import torch
 from torch._six import PY3
 
-try:
-    from sklearn.utils import shuffle
-except ModuleNotFoundError:
+_SKLEARN_AVAILABLE = importlib.util.find_spec("sklearn") is not None
+if _SKLEARN_AVAILABLE:
+    from sklearn.utils import shuffle as sk_shuffle
+else:
     warn('You want to use `sklearn` which is not installed yet,'  # pragma: no-cover
          ' install it with `pip install sklearn`.')
 
@@ -72,8 +74,13 @@ def __init__(
         super(ImageNet, self).__init__(self.split_folder, **kwargs)
         self.root = root
 
+        if not _SKLEARN_AVAILABLE:
+            raise ModuleNotFoundError(  # pragma: no-cover
+                'You want to use `shuffle` function from `scikit-learn` which is not installed yet.'
+            )
+
         # shuffle images first
-        self.imgs = shuffle(self.imgs, random_state=1234)
+        self.imgs = sk_shuffle(self.imgs, random_state=1234)
 
         # partition train set into [train, val]
         if split == 'train':
@@ -98,7 +105,7 @@ def __init__(
         # limit the number of classes
         if num_classes != -1:
             # choose the classes at random (but deterministic)
-            ok_classes = shuffle(list(range(num_classes)), random_state=1234)
+            ok_classes = sk_shuffle(list(range(num_classes)), random_state=1234)
             ok_classes = ok_classes[:num_classes]
             ok_classes = set(ok_classes)
 
@@ -110,7 +117,7 @@ def __init__(
             self.imgs = clean_imgs
 
         # shuffle again for final exit
-        self.imgs = shuffle(self.imgs, random_state=1234)
+        self.imgs = sk_shuffle(self.imgs, random_state=1234)
 
         # list of class_nbs for each image
         idcs = [idx for _, idx in self.imgs]
diff --git a/pl_bolts/datasets/ssl_amdim_datasets.py b/pl_bolts/datasets/ssl_amdim_datasets.py
--- a/pl_bolts/datasets/ssl_amdim_datasets.py
+++ b/pl_bolts/datasets/ssl_amdim_datasets.py
@@ -1,12 +1,14 @@
 from abc import ABC
+import importlib
 from typing import Callable, Optional
 from warnings import warn
 
 import numpy as np
 
-try:
-    from sklearn.utils import shuffle
-except ModuleNotFoundError:
+_SKLEARN_AVAILABLE = importlib.util.find_spec("sklearn") is not None
+if _SKLEARN_AVAILABLE:
+    from sklearn.utils import shuffle as sk_shuffle
+else:
     warn('You want to use `sklearn` which is not installed yet,'  # pragma: no-cover
          ' install it with `pip install sklearn`.')
 
@@ -81,9 +83,14 @@ def select_nb_imgs_per_class(cls, examples, labels, nb_imgs_in_val):
 
     @classmethod
     def deterministic_shuffle(cls, x, y):
+        if not _SKLEARN_AVAILABLE:
+            raise ModuleNotFoundError(  # pragma: no-cover
+                'You want to use `shuffle` function from `scikit-learn` which is not installed yet.'
+            )
+
         n = len(x)
         idxs = list(range(0, n))
-        idxs = shuffle(idxs, random_state=1234)
+        idxs = sk_shuffle(idxs, random_state=1234)
 
         x = x[idxs]
 
diff --git a/pl_bolts/utils/semi_supervised.py b/pl_bolts/utils/semi_supervised.py
--- a/pl_bolts/utils/semi_supervised.py
+++ b/pl_bolts/utils/semi_supervised.py
@@ -1,12 +1,14 @@
+import importlib
 import math
 from warnings import warn
 
 import numpy as np
 import torch
 
-try:
+_SKLEARN_AVAILABLE = importlib.util.find_spec("sklearn") is not None
+if _SKLEARN_AVAILABLE:
     from sklearn.utils import shuffle as sk_shuffle
-except ModuleNotFoundError:
+else:
     warn('You want to use `sklearn` which is not installed yet,'  # pragma: no-cover
          ' install it with `pip install sklearn`.')
 
@@ -36,11 +38,15 @@ def balance_classes(X: np.ndarray, Y: list, batch_size: int):
     Perfect balance
 
     Args:
-
         X: input features
         Y: mixed labels (ints)
         batch_size: the ultimate batch size
     """
+    if not _SKLEARN_AVAILABLE:
+        raise ModuleNotFoundError(  # pragma: no-cover
+            'You want to use `shuffle` function from `scikit-learn` which is not installed yet.'
+        )
+
     nb_classes = len(set(Y))
 
     nb_batches = math.ceil(len(Y) / batch_size)


Test patch: