diff --git a/pl_bolts/callbacks/data_monitor.py b/pl_bolts/callbacks/data_monitor.py index a623951d2e..be5c616de5 100644 --- a/pl_bolts/callbacks/data_monitor.py +++ b/pl_bolts/callbacks/data_monitor.py @@ -4,15 +4,13 @@ import numpy as np import torch import torch.nn as nn -from torch import Tensor -from torch.utils.hooks import RemovableHandle - from pytorch_lightning import Callback from pytorch_lightning import LightningModule, Trainer from pytorch_lightning.loggers import WandbLogger, TensorBoardLogger from pytorch_lightning.utilities import rank_zero_warn from pytorch_lightning.utilities.apply_func import apply_to_collection - +from torch import Tensor +from torch.utils.hooks import RemovableHandle try: import wandb diff --git a/pl_bolts/callbacks/ssl_online.py b/pl_bolts/callbacks/ssl_online.py index 52fb1a796f..788b0f252f 100644 --- a/pl_bolts/callbacks/ssl_online.py +++ b/pl_bolts/callbacks/ssl_online.py @@ -1,11 +1,10 @@ -from typing import List, Optional, Union +from typing import Optional import torch +from pytorch_lightning import Callback from pytorch_lightning.metrics.functional import accuracy from torch.nn import functional as F -from pytorch_lightning import Callback - class SSLOnlineEvaluator(Callback): # pragma: no-cover """ diff --git a/pl_bolts/callbacks/variational.py b/pl_bolts/callbacks/variational.py index b0ce7c7fce..1b945e9631 100644 --- a/pl_bolts/callbacks/variational.py +++ b/pl_bolts/callbacks/variational.py @@ -1,14 +1,14 @@ import math -from warnings import warn import torch from pytorch_lightning.callbacks import Callback +from pl_bolts.utils.warnings import warn_missing_pkg + try: import torchvision except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover class LatentDimInterpolator(Callback): diff --git a/pl_bolts/callbacks/vision/image_generation.py b/pl_bolts/callbacks/vision/image_generation.py index 021b53283d..ac71c0dfd0 100644 --- a/pl_bolts/callbacks/vision/image_generation.py +++ b/pl_bolts/callbacks/vision/image_generation.py @@ -1,13 +1,12 @@ -from warnings import warn - import torch from pytorch_lightning import Callback +from pl_bolts.utils.warnings import warn_missing_pkg + try: import torchvision except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover class TensorboardGenerativeModelImageSampler(Callback): diff --git a/pl_bolts/datamodules/binary_mnist_datamodule.py b/pl_bolts/datamodules/binary_mnist_datamodule.py index 8ef9467acc..fc08c6f5a8 100644 --- a/pl_bolts/datamodules/binary_mnist_datamodule.py +++ b/pl_bolts/datamodules/binary_mnist_datamodule.py @@ -1,16 +1,15 @@ -from warnings import warn - import torch from pytorch_lightning import LightningDataModule from torch.utils.data import DataLoader, random_split +from pl_bolts.utils.warnings import warn_missing_pkg + try: from torchvision import transforms as transform_lib from torchvision.datasets import MNIST from pl_bolts.datasets.mnist_dataset import BinaryMNIST except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover _TORCHVISION_AVAILABLE = False else: _TORCHVISION_AVAILABLE = True diff --git a/pl_bolts/datamodules/cifar10_datamodule.py b/pl_bolts/datamodules/cifar10_datamodule.py index 29713cb68c..19535b2932 100644 --- a/pl_bolts/datamodules/cifar10_datamodule.py +++ b/pl_bolts/datamodules/cifar10_datamodule.py @@ -1,6 +1,5 @@ import os from typing import Optional, Sequence -from warnings import warn import torch from pytorch_lightning import LightningDataModule @@ -8,14 +7,14 @@ from pl_bolts.datasets.cifar10_dataset import TrialCIFAR10 from pl_bolts.transforms.dataset_normalizations import cifar10_normalization +from pl_bolts.utils.warnings import warn_missing_pkg try: from torchvision import transforms as transform_lib from torchvision.datasets import CIFAR10 except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover _TORCHVISION_AVAILABLE = False else: _TORCHVISION_AVAILABLE = True diff --git a/pl_bolts/datamodules/cityscapes_datamodule.py b/pl_bolts/datamodules/cityscapes_datamodule.py index 02d1d223bb..a0d623253e 100644 --- a/pl_bolts/datamodules/cityscapes_datamodule.py +++ b/pl_bolts/datamodules/cityscapes_datamodule.py @@ -1,15 +1,13 @@ -from warnings import warn - -import torch from pytorch_lightning import LightningDataModule -from torch.utils.data import DataLoader, random_split +from torch.utils.data import DataLoader + +from pl_bolts.utils.warnings import warn_missing_pkg try: from torchvision import transforms as transform_lib from torchvision.datasets import Cityscapes except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover _TORCHVISION_AVAILABLE = False else: _TORCHVISION_AVAILABLE = True diff --git a/pl_bolts/datamodules/fashion_mnist_datamodule.py b/pl_bolts/datamodules/fashion_mnist_datamodule.py index 1adc2a84c9..32ccb2ce81 100644 --- a/pl_bolts/datamodules/fashion_mnist_datamodule.py +++ b/pl_bolts/datamodules/fashion_mnist_datamodule.py @@ -1,15 +1,14 @@ -from warnings import warn - import torch from pytorch_lightning import LightningDataModule from torch.utils.data import DataLoader, random_split +from pl_bolts.utils.warnings import warn_missing_pkg + try: from torchvision import transforms as transform_lib from torchvision.datasets import FashionMNIST except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover _TORCHVISION_AVAILABLE = False else: _TORCHVISION_AVAILABLE = True diff --git a/pl_bolts/datamodules/imagenet_datamodule.py b/pl_bolts/datamodules/imagenet_datamodule.py index 806a63ffe1..63ae86d1e0 100644 --- a/pl_bolts/datamodules/imagenet_datamodule.py +++ b/pl_bolts/datamodules/imagenet_datamodule.py @@ -1,17 +1,16 @@ import os from typing import Optional -from warnings import warn from pytorch_lightning import LightningDataModule from torch.utils.data import DataLoader from pl_bolts.transforms.dataset_normalizations import imagenet_normalization +from pl_bolts.utils.warnings import warn_missing_pkg try: from torchvision import transforms as transform_lib except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover _TORCHVISION_AVAILABLE = False else: _TORCHVISION_AVAILABLE = True @@ -19,8 +18,7 @@ try: from pl_bolts.datasets.imagenet_dataset import UnlabeledImagenet except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover _TORCHVISION_AVAILABLE = False else: _TORCHVISION_AVAILABLE = True diff --git a/pl_bolts/datamodules/mnist_datamodule.py b/pl_bolts/datamodules/mnist_datamodule.py index 7229462e97..438914dd5e 100644 --- a/pl_bolts/datamodules/mnist_datamodule.py +++ b/pl_bolts/datamodules/mnist_datamodule.py @@ -1,15 +1,14 @@ -from warnings import warn - import torch from pytorch_lightning import LightningDataModule from torch.utils.data import DataLoader, random_split +from pl_bolts.utils.warnings import warn_missing_pkg + try: from torchvision import transforms as transform_lib from torchvision.datasets import MNIST except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover _TORCHVISION_AVAILABLE = False else: _TORCHVISION_AVAILABLE = True diff --git a/pl_bolts/datamodules/ssl_imagenet_datamodule.py b/pl_bolts/datamodules/ssl_imagenet_datamodule.py index 3ce1f6d6d8..4966363f28 100644 --- a/pl_bolts/datamodules/ssl_imagenet_datamodule.py +++ b/pl_bolts/datamodules/ssl_imagenet_datamodule.py @@ -1,18 +1,17 @@ import os -from warnings import warn from pytorch_lightning import LightningDataModule from torch.utils.data import DataLoader from pl_bolts.datasets.imagenet_dataset import UnlabeledImagenet from pl_bolts.transforms.dataset_normalizations import imagenet_normalization +from pl_bolts.utils.warnings import warn_missing_pkg try: from torchvision import transforms as transform_lib except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover _TORCHVISION_AVAILABLE = False else: _TORCHVISION_AVAILABLE = True diff --git a/pl_bolts/datamodules/stl10_datamodule.py b/pl_bolts/datamodules/stl10_datamodule.py index a6eda48c7e..2ca58b3268 100644 --- a/pl_bolts/datamodules/stl10_datamodule.py +++ b/pl_bolts/datamodules/stl10_datamodule.py @@ -1,6 +1,5 @@ import os from typing import Optional -from warnings import warn import torch from pytorch_lightning import LightningDataModule @@ -8,13 +7,13 @@ from pl_bolts.datasets.concat_dataset import ConcatDataset from pl_bolts.transforms.dataset_normalizations import stl10_normalization +from pl_bolts.utils.warnings import warn_missing_pkg try: from torchvision import transforms as transform_lib from torchvision.datasets import STL10 except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover _TORCHVISION_AVAILABLE = False else: _TORCHVISION_AVAILABLE = True diff --git a/pl_bolts/datamodules/vocdetection_datamodule.py b/pl_bolts/datamodules/vocdetection_datamodule.py index 580aca09f2..008d859c35 100644 --- a/pl_bolts/datamodules/vocdetection_datamodule.py +++ b/pl_bolts/datamodules/vocdetection_datamodule.py @@ -1,16 +1,15 @@ -from warnings import warn - import torch import torchvision.transforms as T from pytorch_lightning import LightningDataModule from torch.utils.data import DataLoader +from pl_bolts.utils.warnings import warn_missing_pkg + try: from torchvision.datasets import VOCDetection except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover _TORCHVISION_AVAILABLE = False else: _TORCHVISION_AVAILABLE = True diff --git a/pl_bolts/datasets/cifar10_dataset.py b/pl_bolts/datasets/cifar10_dataset.py index cd6e50ddef..2b8ef9fc39 100644 --- a/pl_bolts/datasets/cifar10_dataset.py +++ b/pl_bolts/datasets/cifar10_dataset.py @@ -2,16 +2,16 @@ import pickle import tarfile from typing import Tuple, Optional, Sequence, Callable -from warnings import warn import torch from torch import Tensor +from pl_bolts.utils.warnings import warn_missing_pkg + try: from PIL import Image except ModuleNotFoundError: - warn('You want to use `Pillow` which is not installed yet,' # pragma: no-cover - ' install it with `pip install Pillow`.') + warn_missing_pkg('PIL', pypi_name='Pillow') # pragma: no-cover _PIL_AVAILABLE = False else: _PIL_AVAILABLE = True diff --git a/pl_bolts/datasets/mnist_dataset.py b/pl_bolts/datasets/mnist_dataset.py index 0144f96345..cdea04566b 100644 --- a/pl_bolts/datasets/mnist_dataset.py +++ b/pl_bolts/datasets/mnist_dataset.py @@ -1,4 +1,4 @@ -from warnings import warn +from pl_bolts.utils.warnings import warn_missing_pkg try: from torchvision import transforms as transform_lib @@ -11,8 +11,7 @@ try: from PIL import Image except ModuleNotFoundError: - warn('You want to use `Pillow` which is not installed yet,' # pragma: no-cover - ' install it with `pip install Pillow`.') + warn_missing_pkg('PIL', pypi_name='Pillow') # pragma: no-cover _PIL_AVAILABLE = False else: _PIL_AVAILABLE = True diff --git a/pl_bolts/datasets/ssl_amdim_datasets.py b/pl_bolts/datasets/ssl_amdim_datasets.py index 33915564a4..e7cc43a78d 100644 --- a/pl_bolts/datasets/ssl_amdim_datasets.py +++ b/pl_bolts/datasets/ssl_amdim_datasets.py @@ -1,14 +1,14 @@ from abc import ABC from typing import Callable, Optional -from warnings import warn import numpy as np +from pl_bolts.utils.warnings import warn_missing_pkg + try: from torchvision.datasets import CIFAR10 except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover class SSLDatasetMixin(ABC): diff --git a/pl_bolts/models/detection/faster_rcnn.py b/pl_bolts/models/detection/faster_rcnn.py index c6fbb819e8..b6b6674bf8 100644 --- a/pl_bolts/models/detection/faster_rcnn.py +++ b/pl_bolts/models/detection/faster_rcnn.py @@ -1,15 +1,15 @@ from argparse import ArgumentParser -from warnings import warn import pytorch_lightning as pl import torch +from pl_bolts.utils.warnings import warn_missing_pkg + try: from torchvision.models.detection import faster_rcnn, fasterrcnn_resnet50_fpn from torchvision.ops import box_iou except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover def _evaluate_iou(target, pred): diff --git a/pl_bolts/models/mnist_module.py b/pl_bolts/models/mnist_module.py index d9521e1db8..ac024afc44 100644 --- a/pl_bolts/models/mnist_module.py +++ b/pl_bolts/models/mnist_module.py @@ -1,18 +1,18 @@ import os from argparse import ArgumentParser -from warnings import warn import torch from pytorch_lightning import LightningModule, Trainer from torch.nn import functional as F from torch.utils.data import DataLoader, random_split +from pl_bolts.utils.warnings import warn_missing_pkg + try: from torchvision import transforms from torchvision.datasets import MNIST except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover class LitMNIST(LightningModule): diff --git a/pl_bolts/models/rl/common/gym_wrappers.py b/pl_bolts/models/rl/common/gym_wrappers.py index 16bfd2d2e0..1aefbc3231 100644 --- a/pl_bolts/models/rl/common/gym_wrappers.py +++ b/pl_bolts/models/rl/common/gym_wrappers.py @@ -3,18 +3,18 @@ https://github.com/Shmuma/ptan/blob/master/ptan/common/wrappers.py """ import collections -from warnings import warn import gym import gym.spaces import numpy as np import torch +from pl_bolts.utils.warnings import warn_missing_pkg + try: import cv2 except ModuleNotFoundError: - warn('You want to use `openCV` which is not installed yet,' # pragma: no-cover - ' install it with `pip install opencv-python`.') + warn_missing_pkg('cv2', pypi_name='opencv-python') # pragma: no-cover _OPENCV_AVAILABLE = False else: _OPENCV_AVAILABLE = True diff --git a/pl_bolts/models/rl/dqn_model.py b/pl_bolts/models/rl/dqn_model.py index 5b3c7704b6..bb99422c50 100644 --- a/pl_bolts/models/rl/dqn_model.py +++ b/pl_bolts/models/rl/dqn_model.py @@ -5,7 +5,6 @@ import argparse from collections import OrderedDict from typing import Tuple, List, Dict, Optional -from warnings import warn import numpy as np import pytorch_lightning as pl @@ -21,12 +20,12 @@ from pl_bolts.models.rl.common.agents import ValueAgent from pl_bolts.models.rl.common.memory import MultiStepBuffer from pl_bolts.models.rl.common.networks import CNN +from pl_bolts.utils.warnings import warn_missing_pkg try: from pl_bolts.models.rl.common.gym_wrappers import gym, make_environment except ModuleNotFoundError: - warn('You want to use `gym` which is not installed yet,' # pragma: no-cover - ' install it with `pip install gym`.') + warn_missing_pkg('gym') # pragma: no-cover _GYM_AVAILABLE = False else: _GYM_AVAILABLE = True diff --git a/pl_bolts/models/rl/reinforce_model.py b/pl_bolts/models/rl/reinforce_model.py index 9dd2173280..029047d6c6 100644 --- a/pl_bolts/models/rl/reinforce_model.py +++ b/pl_bolts/models/rl/reinforce_model.py @@ -1,7 +1,6 @@ import argparse from collections import OrderedDict from typing import Tuple, List -from warnings import warn import numpy as np import pytorch_lightning as pl @@ -17,11 +16,12 @@ from pl_bolts.datamodules.experience_source import Experience from pl_bolts.models.rl.common.agents import PolicyAgent from pl_bolts.models.rl.common.networks import MLP +from pl_bolts.utils.warnings import warn_missing_pkg try: import gym except ModuleNotFoundError: - warn('You want to use `gym` which is not installed yet, install it with `pip install gym`.') # pragma: no-cover + warn_missing_pkg('gym') # pragma: no-cover _GYM_AVAILABLE = False else: _GYM_AVAILABLE = True diff --git a/pl_bolts/models/rl/vanilla_policy_gradient_model.py b/pl_bolts/models/rl/vanilla_policy_gradient_model.py index e28dcf9c62..568de89983 100644 --- a/pl_bolts/models/rl/vanilla_policy_gradient_model.py +++ b/pl_bolts/models/rl/vanilla_policy_gradient_model.py @@ -1,7 +1,6 @@ import argparse from collections import OrderedDict from typing import Tuple, List -from warnings import warn import numpy as np import pytorch_lightning as pl @@ -16,11 +15,12 @@ from pl_bolts.datamodules import ExperienceSourceDataset from pl_bolts.models.rl.common.agents import PolicyAgent from pl_bolts.models.rl.common.networks import MLP +from pl_bolts.utils.warnings import warn_missing_pkg try: import gym except ModuleNotFoundError: - warn('You want to use `gym` which is not installed yet, install it with `pip install gym`.') # pragma: no-cover + warn_missing_pkg('gym') # pragma: no-cover _GYM_AVAILABLE = False else: _GYM_AVAILABLE = True diff --git a/pl_bolts/models/self_supervised/amdim/datasets.py b/pl_bolts/models/self_supervised/amdim/datasets.py index 6768635acf..d6ad6bd032 100644 --- a/pl_bolts/models/self_supervised/amdim/datasets.py +++ b/pl_bolts/models/self_supervised/amdim/datasets.py @@ -1,16 +1,16 @@ from typing import Optional -from warnings import warn from torch.utils.data import random_split +from pl_bolts.utils.warnings import warn_missing_pkg + try: from torchvision.datasets import STL10 from pl_bolts.datasets.imagenet_dataset import UnlabeledImagenet from pl_bolts.datasets.ssl_amdim_datasets import CIFAR10Mixed from pl_bolts.models.self_supervised.amdim import transforms as amdim_transforms except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover class AMDIMPretraining(): diff --git a/pl_bolts/models/self_supervised/amdim/transforms.py b/pl_bolts/models/self_supervised/amdim/transforms.py index cd30aa3972..f9705d6ae9 100644 --- a/pl_bolts/models/self_supervised/amdim/transforms.py +++ b/pl_bolts/models/self_supervised/amdim/transforms.py @@ -1,12 +1,10 @@ -from warnings import warn - from pl_bolts.transforms.self_supervised import RandomTranslateWithReflect +from pl_bolts.utils.warnings import warn_missing_pkg try: from torchvision import transforms except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover _TORCHVISION_AVAILABLE = False else: _TORCHVISION_AVAILABLE = True diff --git a/pl_bolts/models/self_supervised/cpc/transforms.py b/pl_bolts/models/self_supervised/cpc/transforms.py index 0fa7b10a3f..8d2b8b0d97 100644 --- a/pl_bolts/models/self_supervised/cpc/transforms.py +++ b/pl_bolts/models/self_supervised/cpc/transforms.py @@ -1,12 +1,10 @@ -from warnings import warn - from pl_bolts.transforms.self_supervised import RandomTranslateWithReflect, Patchify +from pl_bolts.utils.warnings import warn_missing_pkg try: from torchvision import transforms except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover _TORCHVISION_AVAILABLE = False else: _TORCHVISION_AVAILABLE = True diff --git a/pl_bolts/models/self_supervised/moco/moco2_module.py b/pl_bolts/models/self_supervised/moco/moco2_module.py index 03b50052b0..ebd3e1583f 100644 --- a/pl_bolts/models/self_supervised/moco/moco2_module.py +++ b/pl_bolts/models/self_supervised/moco/moco2_module.py @@ -12,18 +12,18 @@ from argparse import ArgumentParser from typing import Union -from warnings import warn import pytorch_lightning as pl import torch import torch.nn.functional as F from torch import nn +from pl_bolts.utils.warnings import warn_missing_pkg + try: import torchvision except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover from pl_bolts.metrics import precision_at_k, mean from pl_bolts.models.self_supervised.moco.transforms import ( diff --git a/pl_bolts/models/self_supervised/moco/transforms.py b/pl_bolts/models/self_supervised/moco/transforms.py index adfa88f9be..898654f05b 100644 --- a/pl_bolts/models/self_supervised/moco/transforms.py +++ b/pl_bolts/models/self_supervised/moco/transforms.py @@ -1,14 +1,13 @@ import random -from warnings import warn from pl_bolts.transforms.dataset_normalizations import ( imagenet_normalization, cifar10_normalization, stl10_normalization) +from pl_bolts.utils.warnings import warn_missing_pkg try: from torchvision import transforms except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover _TORCHVISION_AVAILABLE = False else: _TORCHVISION_AVAILABLE = True @@ -16,8 +15,7 @@ try: from PIL import ImageFilter except ModuleNotFoundError: - warn('You want to use `Pillow` which is not installed yet,' # pragma: no-cover - ' install it with `pip install Pillow`.') + warn_missing_pkg('PIL', pypi_name='Pillow') # pragma: no-cover _PIL_AVAILABLE = False else: _PIL_AVAILABLE = True diff --git a/pl_bolts/models/self_supervised/resnets.py b/pl_bolts/models/self_supervised/resnets.py index 4d54b2f009..83c35fc967 100644 --- a/pl_bolts/models/self_supervised/resnets.py +++ b/pl_bolts/models/self_supervised/resnets.py @@ -1,12 +1,11 @@ -from warnings import warn - import torch.nn as nn +from pl_bolts.utils.warnings import warn_missing_pkg + try: from torchvision.models.utils import load_state_dict_from_url except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet50_bn', 'resnet101', 'resnet152', 'resnext50_32x4d', 'resnext101_32x8d', diff --git a/pl_bolts/models/self_supervised/simclr/simclr_module.py b/pl_bolts/models/self_supervised/simclr/simclr_module.py index 92517ed490..607dee9c12 100644 --- a/pl_bolts/models/self_supervised/simclr/simclr_module.py +++ b/pl_bolts/models/self_supervised/simclr/simclr_module.py @@ -1,16 +1,16 @@ from argparse import ArgumentParser -from warnings import warn import pytorch_lightning as pl from torch import nn from torch.nn import functional as F from torch.optim import Adam +from pl_bolts.utils.warnings import warn_missing_pkg + try: from torchvision.models import densenet except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover from pl_bolts.losses.self_supervised_learning import nt_xent_loss from pl_bolts.models.self_supervised.evaluator import Flatten diff --git a/pl_bolts/models/self_supervised/simclr/transforms.py b/pl_bolts/models/self_supervised/simclr/transforms.py index 9a9d2cc12d..a154130768 100644 --- a/pl_bolts/models/self_supervised/simclr/transforms.py +++ b/pl_bolts/models/self_supervised/simclr/transforms.py @@ -1,12 +1,11 @@ -from warnings import warn - import numpy as np +from pl_bolts.utils.warnings import warn_missing_pkg + try: import torchvision.transforms as transforms except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover _TORCHVISION_AVAILABLE = False else: _TORCHVISION_AVAILABLE = True @@ -14,8 +13,7 @@ try: import cv2 except ModuleNotFoundError: - warn('You want to use `opencv-python` which is not installed yet,' # pragma: no-cover - ' install it with `pip install opencv-python`.') + warn_missing_pkg('cv2', pypi_name='opencv-python') # pragma: no-cover class SimCLRTrainDataTransform(object): diff --git a/pl_bolts/models/self_supervised/ssl_finetuner.py b/pl_bolts/models/self_supervised/ssl_finetuner.py index 98cbe224ba..6d67bdea4a 100644 --- a/pl_bolts/models/self_supervised/ssl_finetuner.py +++ b/pl_bolts/models/self_supervised/ssl_finetuner.py @@ -1,11 +1,11 @@ +from typing import Optional, List + import pytorch_lightning as pl -import pytorch_lightning.metrics.functional as plm import torch import torch.nn.functional as F +from pytorch_lightning.metrics import Accuracy from pl_bolts.models.self_supervised import SSLEvaluator -from typing import Optional, List -from pytorch_lightning.metrics import Accuracy class SSLFineTuner(pl.LightningModule): diff --git a/pl_bolts/models/self_supervised/swav/swav_module.py b/pl_bolts/models/self_supervised/swav/swav_module.py index b7cf58ebbf..a9a331ecf4 100644 --- a/pl_bolts/models/self_supervised/swav/swav_module.py +++ b/pl_bolts/models/self_supervised/swav/swav_module.py @@ -4,21 +4,19 @@ import math import os from argparse import ArgumentParser -from warnings import warn +from typing import Callable, Optional import numpy as np import pytorch_lightning as pl import torch import torch.distributed as dist +from pytorch_lightning.utilities import AMPType from torch import nn from torch.optim.optimizer import Optimizer -from typing import Callable, Optional -from pytorch_lightning.utilities import AMPType - from pl_bolts.models.self_supervised.swav.swav_resnet import resnet50, resnet18 -from pl_bolts.transforms.dataset_normalizations import stl10_normalization, cifar10_normalization from pl_bolts.optimizers.lars_scheduling import LARSWrapper +from pl_bolts.transforms.dataset_normalizations import stl10_normalization, cifar10_normalization class SwAV(pl.LightningModule): @@ -474,7 +472,7 @@ def add_model_specific_args(parent_parser): def cli_main(): from pl_bolts.callbacks.ssl_online import SSLOnlineEvaluator from pl_bolts.models.self_supervised.swav.transforms import SwAVTrainDataTransform, SwAVEvalDataTransform - from pl_bolts.datamodules import STL10DataModule, ImagenetDataModule, CIFAR10DataModule + from pl_bolts.datamodules import STL10DataModule, CIFAR10DataModule parser = ArgumentParser() diff --git a/pl_bolts/models/self_supervised/swav/transforms.py b/pl_bolts/models/self_supervised/swav/transforms.py index 2e243ca8c4..de5b72e620 100644 --- a/pl_bolts/models/self_supervised/swav/transforms.py +++ b/pl_bolts/models/self_supervised/swav/transforms.py @@ -1,12 +1,11 @@ -from warnings import warn - import numpy as np +from pl_bolts.utils.warnings import warn_missing_pkg + try: import torchvision.transforms as transforms except ImportError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover _TORCHVISION_AVAILABLE = False else: _TORCHVISION_AVAILABLE = True @@ -14,8 +13,7 @@ try: import cv2 except ImportError: - warn('You want to use `opencv-python` which is not installed yet,' # pragma: no-cover - ' install it with `pip install opencv-python`.') + warn_missing_pkg('cv2', pypi_name='opencv-python') # pragma: no-cover from typing import List diff --git a/pl_bolts/transforms/dataset_normalizations.py b/pl_bolts/transforms/dataset_normalizations.py index 12115c3b1f..c9dae42468 100644 --- a/pl_bolts/transforms/dataset_normalizations.py +++ b/pl_bolts/transforms/dataset_normalizations.py @@ -1,9 +1,9 @@ -from warnings import warn +from pl_bolts.utils.warnings import warn_missing_pkg + try: from torchvision import transforms except ModuleNotFoundError: - warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover - ' install it with `pip install torchvision`.') + warn_missing_pkg('torchvision') # pragma: no-cover def imagenet_normalization(): diff --git a/pl_bolts/transforms/self_supervised/ssl_transforms.py b/pl_bolts/transforms/self_supervised/ssl_transforms.py index 1e51dda4fc..40eff54633 100644 --- a/pl_bolts/transforms/self_supervised/ssl_transforms.py +++ b/pl_bolts/transforms/self_supervised/ssl_transforms.py @@ -1,13 +1,12 @@ -from warnings import warn - import numpy as np import torch.nn.functional as F +from pl_bolts.utils.warnings import warn_missing_pkg + try: from PIL import Image except ModuleNotFoundError: - warn('You want to use `Pillow` which is not installed yet,' # pragma: no-cover - ' install it with `pip install Pillow`.') + warn_missing_pkg('PIL', pypi_name='Pillow') # pragma: no-cover _PIL_AVAILABLE = False else: _PIL_AVAILABLE = True diff --git a/pl_bolts/utils/semi_supervised.py b/pl_bolts/utils/semi_supervised.py index b168b3f7f0..1dd7ab43c1 100644 --- a/pl_bolts/utils/semi_supervised.py +++ b/pl_bolts/utils/semi_supervised.py @@ -1,16 +1,16 @@ import importlib import math -from warnings import warn import numpy as np import torch +from pl_bolts.utils.warnings import warn_missing_pkg + _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`.') + warn_missing_pkg('sklearn', pypi_name='scikit-learn') # pragma: no-cover class Identity(torch.nn.Module): diff --git a/pl_bolts/utils/warnings.py b/pl_bolts/utils/warnings.py new file mode 100644 index 0000000000..9d9dabbc25 --- /dev/null +++ b/pl_bolts/utils/warnings.py @@ -0,0 +1,31 @@ +import os +import warnings + +MISSING_PACKAGE_WARNINGS = {} + + +def warn_missing_pkg(pkg_name: str, pypi_name: str = None, extra_text: str = None, stdout_func=warnings.warn): + """ + Template for warning on missing packages, show them just once. + + Args: + pkg_name: name of missing package + pypi_name: in case that package name differ from PyPI name + extra_text: additional text after the base warning + stdout_func: define used function for streaming warning, use `warnings.warn` or `logging.warning` + + Returns: + number of warning calls + """ + global MISSING_PACKAGE_WARNINGS + if pkg_name not in MISSING_PACKAGE_WARNINGS: + extra_text = os.linesep + extra_text if extra_text else '' + if not pypi_name: + pypi_name = pkg_name + stdout_func(f'You want to use `{pkg_name}` which is not installed yet,' + f' install it with `pip install {pypi_name}`.' + extra_text) + MISSING_PACKAGE_WARNINGS[pkg_name] = 1 + else: + MISSING_PACKAGE_WARNINGS[pkg_name] += 1 + + return MISSING_PACKAGE_WARNINGS[pkg_name]