Repository: pytorch/opacus (language: python)

Issue:
New validation API
Summary: Still WIP

Differential Revision: D26113547




Discussion / hints:
This pull request was **exported** from Phabricator. Differential Revision: [D26113547](https://phabricator.intern.facebook.com/D26113547)
Hi @Darktex! 

Thank you for your pull request. 

We **require** contributors to sign our **Contributor License Agreement**, and yours needs attention.

You currently have a record in our system, but the CLA is no longer valid, and will need to be **resubmitted**.

# Process

In order for us to review and merge your suggested changes, please sign at <https://code.facebook.com/cla>. **If you are contributing on behalf of someone else (eg your employer)**, the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the **pull request will be tagged** with `CLA signed`. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at [cla@fb.com](mailto:cla@fb.com?subject=CLA%20for%20pytorch%2Fopacus%20%23131). Thanks!

Fix patch:
diff --git a/opacus/dp_model_inspector.py b/opacus/dp_model_inspector.py
deleted file mode 100644
index fd95debf..00000000
--- a/opacus/dp_model_inspector.py
+++ /dev/null
@@ -1,222 +0,0 @@
-#!/usr/bin/env python3
-# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
-
-from opacus.grad_sample import GradSampleModule
-from opacus.utils.module_inspection import ModelInspector, get_layer_type
-from torch import nn
-
-
-class IncompatibleModuleException(Exception):
-    r"""
-    Exception class to be thrown in case
-    the given model contains incompatible modules.
-    """
-
-    pass
-
-
-class DPModelInspector:
-    r"""
-    Class to validate if a given module meets the requirements for attaching
-     ``PrivacyEngine``.
-
-    Active checks are listed in the ``DPModelInspector.inspectors`` attribute.
-    """
-
-    def __init__(self, should_throw: bool = True):
-        r"""
-        Args:
-            should_throw: Whether the inspector should throw an exception or
-                return False in case of validation error
-        """
-        self.should_throw = should_throw
-
-        self.inspectors = [
-            # Inspector to check model only consists of sub-modules we support
-            ModelInspector(
-                name="validity",
-                predicate=_is_valid_check,
-                message="grad_sampler method is not yet supported for this module.",
-            ),
-            # Inspector to check for BatchNorms as they could be replaced with groupnorm
-            ModelInspector(
-                name="batchnorm",
-                predicate=_no_batchnorm_check,
-                message="Model contains BatchNorm layers. It is recommended"
-                "That they are replaced with GroupNorm.",
-            ),
-            # Inspector to check that instance norms doesn't track running stats
-            ModelInspector(
-                name="running_stats",
-                predicate=_no_running_stats_instancenorm_check,
-                message="InstanceNorm layer initialised with track_running_stats=True"
-                "This is currently not supported.",
-            ),
-            # Inspector to check the number of groups in Conv2d layers
-            ModelInspector(
-                name="conv_group_number",
-                predicate=_conv_group_number_check,
-                message="Number of groups in Conv2d layer must be either 1 or equal to number of channels.",
-            ),
-            # Inspector to check for LSTM as it can be replaced with DPLSTM
-            ModelInspector(
-                name="lstm",
-                predicate=_no_lstm,
-                message="Model contains LSTM layers. It is recommended that they are"
-                "replaced with DPLSTM.",
-            ),
-            # Inspector to check that the module is in training mode
-            ModelInspector(
-                name="train_mode",
-                predicate=_is_in_training_mode,
-                message=(
-                    "Your model must be in training mode for the PrivacyEngine to compute "
-                    "per-sample gradients. You can put your model in training mode by simply calling "
-                    "module.train(). If you have part of the model that you want to keep frozen, "
-                    "the best approach is to split your model into a frozen backbone and a trainable "
-                    "head and passing only the head to the optimizer and PrivacyEngine. See an example "
-                    "in this colab: https://bit.ly/opacus-dev-day"
-                ),
-            ),
-        ]
-
-    def validate(self, model: nn.Module) -> bool:
-        r"""
-        Runs the validation on the model and all its submodules.
-
-
-        Validation comprises a series of individual
-        :class:`ModelInspectors <opacus.utils.module_inspection.ModelInspector>`,
-        each checking one predicate. Depending on ``should_throw`` flag in
-        the constructor, will either return False or throw
-        :class:`~opacus.dp_model_inspector.IncompatibleModuleException` in case of
-        validation failure.
-
-        Notes:
-            This method is called in :meth:`opacus.privacy_engine.PrivacyEngine.attach`.
-
-        Args:
-            model: The model to validate.
-
-        Returns:
-            True if successful. False if validation fails and ``should_throw == False``
-
-        Raises:
-            IncompatibleModuleException
-                If the validation fails and ``should_throw == True``. Exception message will
-                contain the details of validation failure reason.
-
-        Example:
-            >>> inspector = DPModelInspector()
-            >>> valid_model = nn.Linear(16, 32)
-            >>> is_valid = inspector.validate(valid_model)
-            >>> is_valid
-            True
-            >>> invalid_model = nn.BatchNorm1d(2)
-            >>> is_valid = inspector.validate(invalid_model)
-            # IncompatibleModuleException is thrown.
-        """
-        valid = all(inspector.validate(model) for inspector in self.inspectors)
-        if self.should_throw and not valid:
-            message = "Model contains incompatible modules."
-            for inspector in self.inspectors:
-                if inspector.violators:
-                    message += f"\n{inspector.message}: {inspector.violators}"
-            raise IncompatibleModuleException(message)
-        return valid
-
-
-def _is_valid_check(module: nn.Module) -> bool:
-    r"""
-    Checks if the ``module`` is supported by ``autograd_grad_sample``
-
-    Args:
-        module: The model to validate.
-
-    Returns:
-        True if ``module`` is supported by ``autograd_grad_sample``
-    """
-    return GradSampleModule.is_supported(module)
-
-
-def _is_in_training_mode(module: nn.Module) -> bool:
-    r"""
-    Checks if the ``module`` is in train mode
-
-    Args:
-        module: The model to validate.
-
-    Returns:
-        True if ``module`` is in train mode
-    """
-    return module.training
-
-
-def _conv_group_number_check(module: nn.Module) -> bool:
-    r"""
-    Checks if number of groups in ``nn.Conv2d`` layer is valid
-
-    Args:
-        module: The input module (layer) for which the check occurs.
-
-    Returns:
-        True if number of groups in ``nn.Conv2d`` layer of ``module`` is valid
-    """
-    if isinstance(module, nn.Conv2d):
-        return module.groups == 1 or module.groups == module.in_channels
-
-    return True
-
-
-def _no_batchnorm_check(module: nn.Module) -> bool:
-    r"""
-    Checks if the input module is not BatchNorm.
-
-    This check overlaps with ``_is_valid_check``, but provides more targeted remedy.
-
-    Args:
-        module: The input module
-
-    Returns:
-        True if the input module is not BatchNorm
-    """
-    return not isinstance(module, nn.modules.batchnorm._BatchNorm)
-
-
-def _no_running_stats_instancenorm_check(module: nn.Module) -> bool:
-    r"""
-    Checks that ``InstanceNorm`` layer has ``track_running_stats`` set to False
-
-    Args:
-        module: The input module (layer) for which the check is verified.
-
-    Returns:
-        True if the module is not ``InstanceNorm``, otherwise it returns
-        True if the module (layer) have ``track_running_stats`` set to False,
-        and False otherwise.
-
-    """
-    is_instancenorm = get_layer_type(module) in (
-        "InstanceNorm1d",
-        "InstanceNorm2d",
-        "InstanceNorm3d",
-    )
-
-    if is_instancenorm:
-        return not module.track_running_stats
-    return True
-
-
-def _no_lstm(module: nn.Module) -> bool:
-    r"""
-    Checks if the input module is not LSTM.
-
-    Args:
-        module: The input module
-
-    Returns:
-        True if the input module is not LSTM
-    """
-    is_lstm = True if get_layer_type(module) == "LSTM" else False
-
-    return not is_lstm
diff --git a/opacus/grad_sample/grad_sample_module.py b/opacus/grad_sample/grad_sample_module.py
index c0c9d517..b67e68ba 100644
--- a/opacus/grad_sample/grad_sample_module.py
+++ b/opacus/grad_sample/grad_sample_module.py
@@ -3,16 +3,12 @@
 from __future__ import annotations
 
 from functools import partial
-from typing import Iterable, List, Tuple
+from typing import List, Tuple
 
 import torch
 import torch.nn as nn
 from opacus.layers.dp_lstm import DPLSTM, LSTMLinear
-from opacus.utils.module_inspection import requires_grad
-
-
-class UnsupportedModuleError(ValueError):
-    pass
+from opacus.utils.module_utils import requires_grad, trainable_modules
 
 
 def create_or_accumulate_grad_sample(
@@ -146,7 +142,7 @@ def add_hooks(self, loss_reduction: str = "mean", batch_first: bool = True) -> N
             self._module.autograd_grad_sample_hooks = []
             self.autograd_grad_sample_hooks = self._module.autograd_grad_sample_hooks
 
-        for module in self.trainable_modules():
+        for module in trainable_modules(self._module):
             if type(module) in self.GRAD_SAMPLERS:
                 self.autograd_grad_sample_hooks.append(
                     module.register_forward_hook(self.capture_activations_hook)
@@ -200,28 +196,6 @@ def enable_hooks(self) -> None:
         """
         self.hooks_enabled = True
 
-    def parametrized_modules(self) -> Iterable[nn.Module]:
-        """
-        Recursively iterates over all submodules, returning those that
-        have parameters (as opposed to "wrapper modules" that just organize modules).
-        """
-        yield from (
-            m
-            for m in self._module.modules()
-            if any(p is not None for p in m.parameters(recurse=False))
-        )
-
-    def trainable_modules(self) -> Iterable[nn.Module]:
-        """
-        Recursively iterates over all submodules, returning those that
-        have parameters and are trainable (ie they want a grad).
-        """
-        yield from (
-            m
-            for m in self.parametrized_modules()
-            if any(p.requires_grad for p in m.parameters())
-        )
-
     def __repr__(self):
         return f"GradSampleModule({self._module.__repr__()})"
 
@@ -356,9 +330,28 @@ def rearrange_grad_samples(
 
     @classmethod
     def is_supported(cls, module: nn.Module) -> bool:
-        """Check if this module is supported"""
+        """Check if this individual module is supported"""
         return type(module) in cls.GRAD_SAMPLERS or type(module) is DPLSTM
 
+    @classmethod
+    def validate(
+        cls, module: nn.Module, raise_if_error: bool = False
+    ) -> List[NotImplementedError]:
+        """Validate support for module being wrapped"""
+        errors = []
+        errors.extend(
+            [
+                NotImplementedError(f"grad sampler is not yet implemented for {m}")
+                for m in trainable_modules(module)
+                if not GradSampleModule.is_supported(m)
+            ]
+        )
+        # raise or return errors as needed
+        if raise_if_error and len(errors) > 0:
+            raise NotImplementedError(errors)
+        else:
+            return errors
+
 
 def _get_batch_size(
     module: nn.Module, grad_sample: torch.Tensor, batch_dim: int
diff --git a/opacus/privacy_engine.py b/opacus/privacy_engine.py
index 16117d0b..060a15d3 100644
--- a/opacus/privacy_engine.py
+++ b/opacus/privacy_engine.py
@@ -5,17 +5,17 @@
 from opacus.accountants import RDPAccountant
 from opacus.accountants.rdp import get_noise_multiplier
 from opacus.data_loader import DPDataLoader
+from opacus.distributed import DifferentiallyPrivateDistributedDataParallel as DPDDP
 from opacus.grad_sample.grad_sample_module import GradSampleModule
-from opacus.optimizers import DPOptimizer, DistributedDPOptimizer, DPPerLayerOptimizer, DistributedPerLayerOptimizer
-from opacus.accountants.rdp import get_noise_multiplier
-from opacus.distributed import (  
-    DifferentiallyPrivateDistributedDataParallel     
-        as DPDDP
+from opacus.optimizers import (
+    DPOptimizer,
+    DistributedDPOptimizer,
+    DPPerLayerOptimizer,
+    DistributedPerLayerOptimizer,
 )
-from opacus.data_loader import DPDataLoader
+from opacus.validators.module_validator import ModuleValidator
 from torch import nn, optim
 from torch.utils.data import DataLoader
-from torch.nn.parallel import DistributedDataParallel as DDP
 
 
 def forbid_accumulation_hook(module: nn.Module, _):
@@ -33,6 +33,29 @@ def __init__(self, secure_mode=False):
         self.accountant = RDPAccountant()
         self.secure_mode = secure_mode  # TODO: actually support it
 
+    def is_compatible(
+        self,
+        module: nn.Module,
+        optimizer: Optional[optim.Optimizer],
+        data_loader: Optional[DataLoader],
+    ) -> bool:
+        """
+        Check if task components are compatible with DP.
+        """
+        return ModuleValidator.is_valid(module)
+
+    def validate(
+        self,
+        module: nn.Module,
+        optimizer: Optional[optim.Optimizer],
+        data_loader: Optional[DataLoader],
+    ):
+        """
+        Validate that task components are compatible with DP.
+        Same as ``is_compatible()``, but raises error instead of returning bool.
+        """
+        ModuleValidator.validate(module, raise_if_error=True)
+
     def make_private(
         self,
         module: nn.Module,
@@ -43,15 +66,18 @@ def make_private(
         batch_first: bool = True,
         loss_reduction: str = "mean",
         poisson_sampling: bool = True,
+        try_fix_incompatible_modules: bool = False,
     ):
         distributed = type(module) is DPDDP
-        assert type(module) is not DDP
 
-        # TODO: DP-Specific validation
+        module = self._prepare_model(
+            module, batch_first, loss_reduction, try_fix_incompatible_modules
+        )
         # TODO: either validate consistent dataset or do per-dataset accounting
-        module = self._prepare_model(module, batch_first, loss_reduction)
         if poisson_sampling:
-            data_loader = self._prepare_data_loader(data_loader, distributed=distributed)
+            data_loader = self._prepare_data_loader(
+                data_loader, distributed=distributed
+            )
 
         sample_rate = 1 / len(data_loader)
         expected_batch_size = int(len(data_loader.dataset) * sample_rate)
@@ -67,7 +93,8 @@ def make_private(
 
         def accountant_hook(optim: DPOptimizer):
             # TODO: This works for Poisson for both single-node and distributed
-            # The reason is that the sample rate is the same in both cases (but in distributed mode, each node samples among a subset of the data)
+            # The reason is that the sample rate is the same in both cases (but in
+            # distributed mode, each node samples among a subset of the data)
             self.accountant.step(
                 noise_multiplier=optim.noise_multiplier,
                 sample_rate=sample_rate * optim.accumulated_iterations,
@@ -86,13 +113,21 @@ def make_private_per_layer(
         max_grad_norms: float,
         batch_first: bool = True,
         loss_reduction: str = "mean",
+        poisson_sampling: bool = True,
+        try_fix_incompatible_modules: bool = False,
     ):
-        distributed = type(module) is DDP
-        assert type(module) is not DPDDP
+        distributed = type(module) is DPDDP
+
+        module = self._prepare_model(
+            module, batch_first, loss_reduction, try_fix_incompatible_modules
+        )
 
-        # TODO: DP-Specific validation
         # TODO: either validate consistent dataset or do per-dataset accounting
-        module = self._prepare_model(module, batch_first, loss_reduction)
+        module = self._prepare_model(
+            module,
+            batch_first,
+            loss_reduction,
+        )
         data_loader = self._prepare_data_loader(data_loader, distributed=distributed)
 
         sample_rate = 1 / len(data_loader)
@@ -112,7 +147,8 @@ def make_private_per_layer(
 
         def accountant_hook(optim: DPOptimizer):
             # TODO: This works for Poisson for both single-node and distributed
-            # The reason is that the sample rate is the same in both cases (but in distributed mode, each node samples among a subset of the data)
+            # The reason is that the sample rate is the same in both cases (but in
+            # distributed mode, each node samples among a subset of the data)
             self.accountant.step(
                 noise_multiplier=optim.noise_multiplier,
                 sample_rate=sample_rate * optim.accumulated_iterations,
@@ -137,6 +173,7 @@ def make_private_with_epsilon(
         max_grad_norm: float,
         batch_first: bool = True,
         loss_reduction: str = "mean",
+        try_fix_incompatible_modules: bool = False,
         alphas: Optional[List[float]] = None,
         sigma_min: Optional[float] = None,
         sigma_max: Optional[float] = None,
@@ -159,6 +196,7 @@ def make_private_with_epsilon(
             max_grad_norm=max_grad_norm,
             batch_first=batch_first,
             loss_reduction=loss_reduction,
+            try_fix_incompatible_modules=try_fix_incompatible_modules,
         )
 
     def _prepare_model(
@@ -166,7 +204,14 @@ def _prepare_model(
         module: nn.Module,
         batch_first: bool = True,
         loss_reduction: str = "mean",
+        try_fix_incompatible_modules: bool = False,
     ) -> GradSampleModule:
+        # (fix and) validate
+        if try_fix_incompatible_modules:
+            module = ModuleValidator.fix(module)
+        self.validate(module=module, optimizer=None, data_loader=None)
+
+        # wrap
         if isinstance(module, GradSampleModule):
             return module
         else:
@@ -181,7 +226,7 @@ def _prepare_optimizer(
         max_grad_norm: float,
         expected_batch_size: int,
         loss_reduction: str = "mean",
-        distributed : bool = False,
+        distributed: bool = False,
     ) -> DPOptimizer:
         if isinstance(optimizer, DPOptimizer):
             # TODO: lol rename optimizer optimizer optimizer
@@ -211,7 +256,7 @@ def _prepare_optimizer_per_layer(
         max_grad_norms: float,
         expected_batch_size: int,
         loss_reduction: str = "mean",
-        distributed : bool = False,
+        distributed: bool = False,
     ) -> DPOptimizer:
         if isinstance(optimizer, DPOptimizer):
             # TODO: lol rename optimizer optimizer optimizer
@@ -225,7 +270,9 @@ def _prepare_optimizer_per_layer(
             loss_reduction=loss_reduction,
         )
 
-    def _prepare_data_loader(self, data_loader: DataLoader, distributed: bool) -> DataLoader:
+    def _prepare_data_loader(
+        self, data_loader: DataLoader, distributed: bool
+    ) -> DataLoader:
         return DPDataLoader.from_data_loader(data_loader, distributed=distributed)
 
     # TODO: default delta value?
diff --git a/opacus/utils/module_inspection.py b/opacus/utils/module_inspection.py
deleted file mode 100644
index 44504673..00000000
--- a/opacus/utils/module_inspection.py
+++ /dev/null
@@ -1,121 +0,0 @@
-#!/usr/bin/env python3
-# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
-r"""
-This module includes utils for inspecting model layers using specified
-predicates to check for conditions, getting layer type etc.
-"""
-from typing import Callable, Optional
-
-from torch import nn
-
-
-class ModelInspector:
-    """
-    An inspector of models given a specific predicate. If a module
-    has children the predicate is checked on all children recursively.
-
-    Example:
-        >>>  inspector = ModelInspector('simple', lambda x: isinstance(x, Conv2d))
-        >>>  print(inspector.validate(nn.Conv2d(1, 1, 1)))
-        True
-    """
-
-    def __init__(
-        self,
-        name: str,
-        predicate: Callable[[nn.Module], bool],
-        check_leaf_nodes_only: bool = True,
-        message: Optional[str] = None,
-    ):
-        """
-        Args:
-            name: String to represent the predicate.
-            predicate: Callable boolean function which tests a hypothesis on a module.
-            check_leaf_nodes_only: Flag to check only leaf nodes of a module. Here
-                leaf nodes are the ones that have parameters of their own.
-            message: Optional value to hold a message about violating this predicate.
-
-        Notes:
-            The predicates will not be applied on non-leaf modules unless
-            ``check_leaf_nodes_only`` is set to False. E.g. A predicate like:
-
-            ``lambda model: isinstance(model, nn.Sequential)``
-
-            will always return True unless ``check_leaf_nodes_only`` is set.
-        """
-        self.name = name
-        if check_leaf_nodes_only:
-            self.predicate = (
-                lambda x: has_no_param(x) or not requires_grad(x) or predicate(x)
-            )
-        else:
-            self.predicate = predicate
-        self.message = message
-        self.violators = []
-        # List that contains the module names that have violated the
-        # predicate. The list does not get automatically emptied if
-        # the predicate is applied on multiple modules.
-
-    def validate(self, model: nn.Module) -> bool:
-        """
-        Checks if the provided module satisfies the predicate specified
-        upon creation of the :class:`~opacus.utils.ModelInspector`.
-
-        Args:
-            model: PyTorch module on which the predicate must be evaluated
-                and satisfied.
-
-        Returns:
-            Flag indicate if predicate is satisfied.
-        """
-        valid = True
-        for name, module in model.named_modules():
-            if not self.predicate(module):
-                valid = False
-                self.violators.append(f"{name} ({get_layer_type(module)})")
-        return valid
-
-
-def has_no_param(module: nn.Module) -> bool:
-    """
-    Checks if a module does not have any parameters.
-
-    Args:
-        module: The module on which this function is being evaluated.
-
-    Returns:
-        Flag indicating if the provided module does not have any
-        parameters.
-    """
-    has_params = any(p is not None for p in module.parameters(recurse=False))
-    return not has_params
-
-
-def requires_grad(module: nn.Module, recurse: bool = False) -> bool:
-    """
-    Checks if any parameters in a specified module require gradients.
-
-    Args:
-        module: PyTorch module whose parameters are examined
-        recurse: Flag specifying if the gradient requirement check should
-            be applied recursively to sub-modules of the specified module
-
-    Returns:
-        Flag indicate if any parameters require gradients
-    """
-    requires_grad = any(p.requires_grad for p in module.parameters(recurse))
-    return requires_grad
-
-
-def get_layer_type(layer: nn.Module) -> str:
-    """
-    Returns the name of the type of the given layer.
-
-    Args:
-        layer: The module corresponding to the layer whose type
-            is being queried.
-
-    Returns:
-        Name of the class of the layer
-    """
-    return layer.__class__.__name__
diff --git a/opacus/utils/module_modification.py b/opacus/utils/module_modification.py
deleted file mode 100644
index 401ea6d7..00000000
--- a/opacus/utils/module_modification.py
+++ /dev/null
@@ -1,170 +0,0 @@
-#!/usr/bin/env python3
-# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
-r"""
-This module includes utils for modifying model layers, replacing layers etc.
-"""
-from typing import Callable, Type
-
-from torch import nn
-
-
-def _replace_child(
-    root: nn.Module, child_name: str, converter: Callable[[nn.Module], nn.Module]
-) -> None:
-    """
-    Converts a sub-module to a new module given a helper
-    function, the root module and a string representing
-    the name of the submodule to be replaced.
-
-    Args:
-        root: Root module whose sub module must be replaced.
-        child_name: Name of submodule that must be replaced.
-        converter: Function or a lambda that takes a module
-            (the submodule to be replaced) and returns its
-            replacement.
-    """
-    # find the immediate parent
-    parent = root
-    nameList = child_name.split(".")
-    for name in nameList[:-1]:
-        parent = parent._modules[name]
-    # set to identity
-    parent._modules[nameList[-1]] = converter(parent._modules[nameList[-1]])
-
-
-def replace_all_modules(
-    root: nn.Module,
-    target_class: Type[nn.Module],
-    converter: Callable[[nn.Module], nn.Module],
-) -> nn.Module:
-    """
-    Converts all the submodules (of root) that have the same
-    type as target_class, given a converter, a module root,
-    and a target class type.
-
-    This method is useful for replacing modules that are not
-    supported by the Privacy Engine.
-
-    Args:
-        root: Model instance, potentially with sub-modules
-        target_class: Target class that needs to be replaced.
-        converter: Function or a lambda that converts an instance
-            of a given target_class to another nn.Module.
-
-    Returns:
-        Module with all the target_class types replaced using the
-        converter. root is modified and is equal to the return value.
-
-    Example:
-        >>>  from torchvision.models import resnet18
-        >>>  from torch import nn
-        >>>  model = resnet18()
-        >>>  print(model.layer1[0].bn1)
-        BatchNorm2d(64, eps=1e-05, ...
-        >>>  model = replace_all_modules(model, nn.BatchNorm2d, lambda _: nn.Identity())
-        >>>  print(model.layer1[0].bn1)
-        Identity()
-    """
-    # base case
-    if isinstance(root, target_class):
-        return converter(root)
-
-    for name, obj in root.named_modules():
-        if isinstance(obj, target_class):
-            _replace_child(root, name, converter)
-    return root
-
-
-def _batchnorm_to_instancenorm(module: nn.modules.batchnorm._BatchNorm) -> nn.Module:
-    """
-    Converts a BatchNorm module to the corresponding InstanceNorm module
-
-    Args:
-        module: BatchNorm module to be replaced
-
-    Returns:
-        InstanceNorm module that can replace the BatchNorm module provided
-    """
-
-    def matchDim():
-        if isinstance(module, nn.BatchNorm1d):
-            return nn.InstanceNorm1d
-        elif isinstance(module, nn.BatchNorm2d):
-            return nn.InstanceNorm2d
-        elif isinstance(module, nn.BatchNorm3d):
-            return nn.InstanceNorm3d
-
-    return matchDim()(module.num_features)
-
-
-def _batchnorm_to_groupnorm(module: nn.modules.batchnorm._BatchNorm) -> nn.Module:
-    """
-    Converts a BatchNorm ``module`` to GroupNorm module.
-    This is a helper function.
-
-    Args:
-        module: BatchNorm module to be replaced
-
-    Returns:
-        GroupNorm module that can replace the BatchNorm module provided
-
-    Notes:
-        A default value of 32 is chosen for the number of groups based on the
-        paper *Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour*
-        https://arxiv.org/pdf/1706.02677.pdf
-    """
-    return nn.GroupNorm(min(32, module.num_features), module.num_features, affine=True)
-
-
-def nullify_batchnorm_modules(root: nn.Module) -> nn.Module:
-    """
-    Replaces all the BatchNorm submodules (e.g. :class:`torch.nn.BatchNorm1d`,
-    :class:`torch.nn.BatchNorm2d` etc.) in ``root`` with :class:`torch.nn.Identity`.
-
-    Args:
-        root: Module for which to replace BatchNorm submodules.
-
-    Returns:
-        Module with all the BatchNorm sub modules replaced with
-        Identity. ``root`` is modified and is equal to the return value.
-
-    Notes:
-        Most of the times replacing a BatchNorm module with Identity
-        will heavily affect convergence of the model.
-    """
-    return replace_all_modules(
-        root, nn.modules.batchnorm._BatchNorm, lambda _: nn.Identity()
-    )
-
-
-def convert_batchnorm_modules(
-    model: nn.Module,
-    converter: Callable[
-        [nn.modules.batchnorm._BatchNorm], nn.Module
-    ] = _batchnorm_to_groupnorm,
-) -> nn.Module:
-    """
-    Converts all BatchNorm modules to another module
-    (defaults to GroupNorm) that is privacy compliant.
-
-    Args:
-        model: Module instance, potentially with sub-modules
-        converter: Function or a lambda that converts an instance of a
-            Batchnorm to another nn.Module.
-
-    Returns:
-        Model with all the BatchNorm types replaced by another operation
-        by using the provided converter, defaulting to GroupNorm if one
-        isn't provided.
-
-    Example:
-        >>>  from torchvision.models import resnet50
-        >>>  from torch import nn
-        >>>  model = resnet50()
-        >>>  print(model.layer1[0].bn1)
-        BatchNorm2d module details
-        >>>  model = convert_batchnorm_modules(model)
-        >>>  print(model.layer1[0].bn1)
-        GroupNorm module details
-    """
-    return replace_all_modules(model, nn.modules.batchnorm._BatchNorm, converter)
diff --git a/opacus/utils/module_utils.py b/opacus/utils/module_utils.py
new file mode 100644
index 00000000..2a28569a
--- /dev/null
+++ b/opacus/utils/module_utils.py
@@ -0,0 +1,94 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+import io
+import logging
+import sys
+from collections import OrderedDict
+from typing import Iterable
+
+import torch
+import torch.nn as nn
+
+logging.basicConfig(
+    format="%(asctime)s:%(levelname)s:%(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+    stream=sys.stderr,
+)
+logger = logging.getLogger(__name__)
+logger.setLevel(level=logging.INFO)
+
+
+def parametrized_modules(module: nn.Module) -> Iterable[nn.Module]:
+    """
+    Recursively iterates over all submodules, returning those that
+    have parameters (as opposed to "wrapper modules" that just organize modules).
+    """
+    yield from (
+        m
+        for m in module.modules()
+        if any(p is not None for p in m.parameters(recurse=False))
+    )
+
+
+def trainable_modules(module: nn.Module) -> Iterable[nn.Module]:
+    """
+    Recursively iterates over all submodules, returning those that
+    have parameters and are trainable (ie they want a grad).
+    """
+    yield from (
+        m
+        for m in parametrized_modules(module)
+        if any(p.requires_grad for p in m.parameters(recurse=False))
+    )
+
+
+def requires_grad(module: nn.Module, recurse: bool = False) -> bool:
+    """
+    Checks if any parameters in a specified module require gradients.
+
+    Args:
+        module: PyTorch module whose parameters are to be examined.
+        recurse: Flag specifying if the gradient requirement check should
+            be applied recursively to sub-modules of the specified module
+
+    Returns:
+        Flag indicate if any parameters require gradients
+    """
+    requires_grad = any(p.requires_grad for p in module.parameters(recurse))
+    return requires_grad
+
+
+def clone_module(module: nn.Module) -> nn.Module:
+    """
+    Handy utility to clone an nn.Module. PyTorch doesn't always support copy.deepcopy(), so it is
+    just easier to serialize the model to a BytesIO and read it from there.
+
+    Args:
+        module: The module to clone
+
+    Returns:
+        The clone of ``module``
+    """
+    with io.BytesIO() as bytesio:
+        torch.save(module, bytesio)
+        bytesio.seek(0)
+        module_copy = torch.load(bytesio)
+    return module_copy
+
+
+def are_state_dict_equal(sd1: OrderedDict, sd2: OrderedDict):
+    if len(sd1) != len(sd2):
+        logger.error(f"Length mismatch: {len(sd1)} vs {len(sd2)}")
+        return False
+
+    for k1, v1 in sd1.items():
+        # check that all keys are accounted for.
+        if k1 not in sd2:
+            logger.error(f"Key missing: {k1} not in {sd2}")
+            return False
+        # check that value tensors are equal.
+        v2 = sd2[k1]
+        if not torch.allclose(v1, v2):
+            logger.error(f"Tensor mismatch: {v1} vs {v2}")
+            return False
+    return True
diff --git a/opacus/validators/__init__.py b/opacus/validators/__init__.py
new file mode 100644
index 00000000..43afd27b
--- /dev/null
+++ b/opacus/validators/__init__.py
@@ -0,0 +1,17 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+from .batch_norm import validate, fix  # noqa
+from .ddp import validate, fix  # noqa
+from .instance_norm import validate, fix  # noqa
+from .lstm import validate, fix  # noqa
+from .module_validator import ModuleValidator
+from .multihead_attention import validate, fix  # noqa
+from .utils import register_module_validator, register_module_fixer
+
+
+__all__ = [
+    "ModuleValidator",
+    "register_module_validator",
+    "register_module_fixer",
+]
diff --git a/opacus/validators/batch_norm.py b/opacus/validators/batch_norm.py
new file mode 100644
index 00000000..738b9fd9
--- /dev/null
+++ b/opacus/validators/batch_norm.py
@@ -0,0 +1,118 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+import logging
+import sys
+from typing import Union
+
+import torch.nn as nn
+
+from .errors import ShouldReplaceModuleError, UnsupportableModuleError
+from .utils import register_module_validator, register_module_fixer
+
+logging.basicConfig(
+    format="%(asctime)s:%(levelname)s:%(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+    stream=sys.stderr,
+)
+logger = logging.getLogger(__name__)
+logger.setLevel(level=logging.INFO)
+
+
+BATCHNORM = Union[nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm]
+INSTANCENORM = Union[nn.InstanceNorm1d, nn.InstanceNorm2d, nn.InstanceNorm3d]
+
+
+@register_module_validator(
+    [nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm]
+)
+def validate(module: BATCHNORM) -> None:
+    return [
+        ShouldReplaceModuleError(
+            "BatchNorm cannot support training with differential privacy. "
+            "The reason for it is that BatchNorm makes each sample's normalized value "
+            "depend on its peers in a batch, ie the same sample x will get normalized to "
+            "a different value depending on who else is on its batch. "
+            "Privacy-wise, this means that we would have to put a privacy mechanism there too. "
+            "While it can in principle be done, there are now multiple normalization layers that "
+            "do not have this issue: LayerNorm, InstanceNorm and their generalization GroupNorm "
+            "are all privacy-safe since they don't have this property."
+            "We offer utilities to automatically replace BatchNorms to GroupNorms and we will "
+            "release pretrained models to help transition, such as GN-ResNet ie a ResNet using "
+            "GroupNorm, pretrained on ImageNet"
+        )
+    ]
+
+
+@register_module_fixer(
+    [nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm]
+)
+def fix(module: BATCHNORM) -> nn.GroupNorm:
+    logger.info("The default fixer replaces BatchNorm with GroupNorm."
+    " The batch_norm validator module also offeers implementations to replace"
+    " it with InstanceNorm or Identity. Please check them out and override the"
+    " fixer incase those are more suitable for your needs.")
+    return _batchnorm_to_groupnorm(module)
+
+
+def _batchnorm_to_groupnorm(module: BATCHNORM) -> nn.GroupNorm:
+    """
+    Converts a BatchNorm ``module`` to GroupNorm module.
+    This is a helper function.
+
+    Args:
+        module: BatchNorm module to be replaced
+
+    Returns:
+        GroupNorm module that can replace the BatchNorm module provided
+
+    Notes:
+        A default value of 32 is chosen for the number of groups based on the
+        paper *Group Normalization* https://arxiv.org/abs/1803.08494
+    """
+    return nn.GroupNorm(
+        min(32, module.num_features), module.num_features, affine=module.affine
+    )
+
+
+def _batchnorm_to_instancenorm(module: BATCHNORM) -> INSTANCENORM:
+    """
+    Converts a BatchNorm module to the corresponding InstanceNorm module
+
+    Args:
+        module: BatchNorm module to be replaced
+
+    Returns:
+        InstanceNorm module that can replace the BatchNorm module provided
+    """
+
+    def match_dim():
+        if isinstance(module, nn.BatchNorm1d):
+            return nn.InstanceNorm1d
+        elif isinstance(module, nn.BatchNorm2d):
+            return nn.InstanceNorm2d
+        elif isinstance(module, nn.BatchNorm3d):
+            return nn.InstanceNorm3d
+        elif isinstance(module, nn.SyncBatchNorm):
+            raise UnsupportableModuleError(
+                "There is no equivalent InstanceNorm module to replace"
+                " SyncBatchNorm with. Consider replacing it with GroupNorm instead."
+            )
+
+    return match_dim()(module.num_features)
+
+
+def _nullify_batch_norm():
+    """
+    Replaces all the BatchNorm with :class:`torch.nn.Identity`.
+    Args:
+        module: BatchNorm module to be replaced
+
+    Returns:
+        InstanceNorm module that can replace the BatchNorm module provided
+
+    Notes:
+        Most of the times replacing a BatchNorm module with Identity
+        will heavily affect convergence of the model.
+    """
+    return nn.Identity()
diff --git a/opacus/validators/ddp.py b/opacus/validators/ddp.py
new file mode 100644
index 00000000..ad0c928e
--- /dev/null
+++ b/opacus/validators/ddp.py
@@ -0,0 +1,32 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+from .errors import ShouldReplaceModuleError
+from .utils import register_module_validator, register_module_fixer
+
+
+@register_module_validator(DDP)
+def validate(module: DDP) -> None:
+    return [
+        ShouldReplaceModuleError(
+            "We do not support DistributedDataParallel as we need to perform"
+            " additonal synchronization across workers. We provide implementation"
+            " DifferentiallyPrivateDistributedDataParallel, which is compatible"
+            " with our DP. Please use the recommended replacement!"
+        )
+    ]
+
+
+@register_module_fixer(DDP)
+def fix(module: DDP) -> None:
+    raise ShouldReplaceModuleError(
+        "We do not support DistributedDataParallel as we need to perform"
+        " additonal synchronization across workers. We provide implementation"
+        " DifferentiallyPrivateDistributedDataParallel, which is compatible"
+        " with our DP. Unfortunately, automatic fix for DDP module is not possible"
+        " since torch distributed process group initialization infomation is "
+        " unavailable to PrivacyEngine. Please fix this manually."
+    )
diff --git a/opacus/validators/errors.py b/opacus/validators/errors.py
new file mode 100644
index 00000000..81cee45d
--- /dev/null
+++ b/opacus/validators/errors.py
@@ -0,0 +1,55 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+class UnsupportedError(ValueError):
+    """
+    Raised for unsupported components in Opacus.
+    """
+
+    pass
+
+
+class UnsupportedModuleError(UnsupportedError):
+    """
+    Raised for unsupported modules in Opacus.
+    """
+
+
+class UnsupportableModuleError(UnsupportedModuleError):
+    """
+    Raised whenever there is a module we can't support ever.
+    BatchNorm is the largest offender.
+    """
+
+    pass
+
+
+class NotYetSupportedModuleError(UnsupportedModuleError):
+    """
+    Raised whenever there is a module that we don't yet support.
+    This is the "catch-all": the number of modules we won't ever support
+    is very short, so a priori if we don't support it now it doesn't mean
+    we can't extend support later (and PRs are welcome!!).
+    """
+
+    pass
+
+
+class ShouldReplaceModuleError(UnsupportedModuleError):
+    """
+    Raised whenever there is a module that we don't support as-is but we do support via
+    replacement (and we have made a replacement ourselves).
+    """
+
+    pass
+
+
+class IllegalModuleConfigurationError(UnsupportedModuleError):
+    """
+    Raised whenever there is an illegal configuration of the training setup.
+    Examples include, not setting the module to training mode, enabling running stats
+    in InstanceNorm, etc.
+    """
+
+    pass
diff --git a/opacus/validators/instance_norm.py b/opacus/validators/instance_norm.py
new file mode 100644
index 00000000..f08fe2f6
--- /dev/null
+++ b/opacus/validators/instance_norm.py
@@ -0,0 +1,39 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+from typing import Union
+
+import torch.nn as nn
+from opacus.utils.module_utils import clone_module
+
+from .errors import IllegalModuleConfigurationError
+from .utils import register_module_validator, register_module_fixer
+
+
+INSTANCENORM = Union[nn.InstanceNorm1d, nn.InstanceNorm2d, nn.InstanceNorm3d]
+
+
+@register_module_validator([nn.InstanceNorm1d, nn.InstanceNorm2d, nn.InstanceNorm3d])
+def validate(module: INSTANCENORM) -> None:
+    return (
+        [
+            IllegalModuleConfigurationError(
+                "We do not support tracking running stats with differential privacy. "
+                "To support it, we would have to add a DP mechanism for these statistics too, "
+                "which would incur a privacy cost for little value in model accuracy. "
+                "Just say no to running stats :)"
+            )
+        ]
+        if module.track_running_stats
+        else []
+    )
+
+
+@register_module_fixer([nn.InstanceNorm1d, nn.InstanceNorm2d, nn.InstanceNorm3d])
+def fix(module: INSTANCENORM) -> INSTANCENORM:
+    if len(validate(module)) == 0:
+        return module
+    # else
+    new_module = clone_module(module)
+    new_module.track_running_stats = False
+    return new_module
diff --git a/opacus/validators/lstm.py b/opacus/validators/lstm.py
new file mode 100644
index 00000000..99b770d6
--- /dev/null
+++ b/opacus/validators/lstm.py
@@ -0,0 +1,35 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+import torch.nn as nn
+from opacus.layers import DPLSTM
+
+from .errors import ShouldReplaceModuleError
+from .utils import register_module_validator, register_module_fixer
+
+
+@register_module_validator(nn.LSTM)
+def validate(module: nn.LSTM) -> None:
+    return [
+        ShouldReplaceModuleError(
+            "We do not support nn.LSTM because its implementation uses special "
+            "modules. We have written a DPLSTM class that is a drop-in replacement "
+            "which is compatible with our Grad Sample hooks. Please run the recommended "
+            "replacement!"
+        )
+    ]
+
+
+@register_module_fixer(nn.LSTM)
+def fix(module: nn.LSTM) -> DPLSTM:
+    dplstm = DPLSTM(
+        input_size=module.input_size,
+        hidden_size=module.hidden_size,
+        num_layers=module.num_layers,
+        bias=module.bias,
+        batch_first=module.batch_first,
+        dropout=module.dropout,
+        bidirectional=module.bidirectional,
+    )
+    dplstm.load_state_dict(module.state_dict())
+    return dplstm
diff --git a/opacus/validators/module_validator.py b/opacus/validators/module_validator.py
new file mode 100644
index 00000000..71d8405f
--- /dev/null
+++ b/opacus/validators/module_validator.py
@@ -0,0 +1,158 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+import logging
+import sys
+from typing import List
+
+import torch.nn as nn
+from opacus.grad_sample.grad_sample_module import GradSampleModule
+from opacus.utils.module_utils import clone_module
+from opacus.validators.errors import (
+    IllegalModuleConfigurationError,
+    UnsupportedModuleError,
+)
+
+logging.basicConfig(
+    format="%(asctime)s:%(levelname)s:%(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+    stream=sys.stderr,
+)
+logger = logging.getLogger(__name__)
+logger.setLevel(level=logging.INFO)
+
+
+class ModuleValidator:
+    """
+    Encapsulates all the validation logic required by Opacus.
+    Also works as a namespace to hold registered validators and fixers.
+    """
+
+    VALIDATORS = {}
+    FIXERS = {}
+
+    @classmethod
+    def validate(
+        cls, module: nn.Module, raise_if_error: bool = False
+    ) -> List[UnsupportedModuleError]:
+        """
+        Validate module and sub_modules by running registered custom validators.
+        Returns or raises excpetions depending on ``raise_if_error`` flag.
+
+        Args:
+            module: The root module to validate.
+            raise_if_error: Boolean to indicate whether to raise errors or return
+            the list of errors.
+
+        Raises:
+            UnsupportedModuleError in case of validation failures.
+        """
+        errors = []
+        # 1. validate that module is in training mode
+        if not module.training:
+            errors.append(
+                IllegalModuleConfigurationError("Model needs to be in training mode")
+            )
+        # 2. validate that all trainable modules are supported by GradSampleModule.
+        errors.extend(GradSampleModule.validate(module=module, raise_if_error=False))
+        # 3. perform module specific validations.
+        for _, sub_module in module.named_modules():
+            if type(sub_module) in ModuleValidator.VALIDATORS:
+                sub_module_validator = ModuleValidator.VALIDATORS[type(sub_module)]
+                errors.extend(sub_module_validator(sub_module))
+        # raise/return as needed
+        if raise_if_error and len(errors) > 0:
+            raise UnsupportedModuleError(errors)
+        else:
+            return errors
+
+    @classmethod
+    def is_valid(cls, module: nn.Module) -> bool:
+        """
+        Check if module and sub_modules are valid by running registered custom validators.
+
+        Args:
+            module: The root module to validate.
+
+        Returns:
+            bool
+        """
+        return len(cls.validate(module, raise_if_error=False)) == 0
+
+    @classmethod
+    def fix(cls, module: nn.Module) -> nn.Module:
+        """
+        Make the module and sub_modules DP compatible by running registered custom fixers.
+
+        Args:
+            module: The root module to be made compatible.
+
+        Returns:
+            Fixed module.
+        """
+        module = clone_module(module)
+        # iterate over all sub_modules
+        #   Need to get sub_module names first as we will be changing
+        #   inside the the loop.
+        sub_module_names = [name for name, _ in module.named_modules()]
+        for sub_module_name in sub_module_names:
+            # get sub_module
+            sub_module = module.get_submodule(sub_module_name)
+            # if sub_module has a registered fixer
+            if type(sub_module) in ModuleValidator.FIXERS:
+                # get a repalcement for sub_module
+                sub_module_fixer = ModuleValidator.FIXERS[type(sub_module)]
+                new_sub_module = sub_module_fixer(sub_module)
+                # get module after replacement.
+                module = cls._repalce_sub_module(
+                    root=module,
+                    sub_module_name=sub_module_name,
+                    new_sub_module=new_sub_module,
+                )
+                # log it
+                logger.info(
+                    f"Replaced sub_module {sub_module_name} : {sub_module}"
+                    f" with {new_sub_module}"
+                )
+        # return fixed module
+        return module
+
+    @classmethod
+    def _repalce_sub_module(
+        cls,
+        root: nn.Module,
+        sub_module_name: str,
+        new_sub_module: nn.Module,
+    ) -> None:
+        sub_module_path = sub_module_name.split(".")
+        if (
+            len(sub_module_path) == 1 and sub_module_path[0] == ""
+        ):  # root is the only sub_module of root
+            return new_sub_module
+        else:  # repalce root's descendant
+            sub_module_parent = root
+            for name in sub_module_path[:-1]:  # descend down to sub_module
+                sub_module_parent = sub_module_parent._modules[name]
+            sub_module_parent._modules[sub_module_path[-1]] = new_sub_module
+        return root
+
+    @classmethod
+    def fix_and_validate(cls, module: nn.Module) -> nn.Module:
+        """
+        Fix the module and sub_modules first, and then run validation.
+
+        Args:
+            module: The root module to be fixed and validted
+
+        Returns:
+            Fixed module.
+
+        Raises:
+            UnsupportedModuleError in case of validation failures.
+        """
+        # 1. replace any fixable modules
+        fixed_module = cls.fix(module)
+        # 2. perform module specific validations.
+        cls.validate(fixed_module, raise_if_error=True)
+        # return fixed module
+        return fixed_module
diff --git a/opacus/validators/multihead_attention.py b/opacus/validators/multihead_attention.py
new file mode 100644
index 00000000..fc159b42
--- /dev/null
+++ b/opacus/validators/multihead_attention.py
@@ -0,0 +1,36 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+import torch.nn as nn
+from opacus.layers import DPMultiheadAttention
+
+from .errors import ShouldReplaceModuleError
+from .utils import register_module_validator, register_module_fixer
+
+
+@register_module_validator(nn.MultiheadAttention)
+def validate(module: nn.MultiheadAttention) -> None:
+    return [
+        ShouldReplaceModuleError(
+            "We do not support nn.MultiheadAttention because its implementation uses special "
+            "modules. We have written a DPMultiheadAttention class that is a drop-in replacement "
+            "which is compatible with our Grad Sample hooks. Please run the recommended "
+            "replacement!"
+        )
+    ]
+
+
+@register_module_fixer(nn.MultiheadAttention)
+def fix(module: nn.MultiheadAttention) -> DPMultiheadAttention:
+    dp_attn = DPMultiheadAttention(
+        embed_dim=module.embed_dim,
+        num_heads=module.num_heads,
+        dropout=module.dropout,
+        bias=module.in_proj_bias is not None,
+        add_bias_kv=module.bias_k is not None,
+        add_zero_attn=module.add_zero_attn,
+        kdim=module.kdim,
+        vdim=module.vdim,
+    )
+    dp_attn.load_state_dict(module.state_dict())
+    return dp_attn
diff --git a/opacus/validators/utils.py b/opacus/validators/utils.py
new file mode 100644
index 00000000..738d853d
--- /dev/null
+++ b/opacus/validators/utils.py
@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+from typing import Sequence, Union
+
+from .module_validator import ModuleValidator
+
+DEFAULT_MODULE_VALIDATOR = ModuleValidator
+
+
+def register_module_validator(
+    target_class_or_classes: Union[type, Sequence[type]],
+    validator_class: type = DEFAULT_MODULE_VALIDATOR,
+):
+    """
+    Registers the decorated function as the ``validator`` of ``target_class_or_classes``, which is
+    the function that will be invoked every time you want to validate that a module is compatible
+    for training with Opacus.
+    You may supply your own validator_class that holds the registry of VALIDATORS.
+    The signature of every validator is always the same:
+
+    >>> @register_module_validator(nn.MyCustomClass)
+    >>> def validate(module: nn.Module, **kwargs) -> List[opacus.validators.errors.UnsupportedError]:
+    >>>    pass
+
+    It may help you to take a look at the existing validator inside Opacus, under ``opacus.validators.``
+    """
+
+    def decorator(f):
+        target_classes = (
+            target_class_or_classes
+            if isinstance(target_class_or_classes, Sequence)
+            else [target_class_or_classes]
+        )
+        for target_class in target_classes:
+            validator_class.VALIDATORS[target_class] = f
+        return f
+
+    return decorator
+
+
+def register_module_fixer(
+    target_class_or_classes: Union[type, Sequence[type]],
+    validator_class: type = DEFAULT_MODULE_VALIDATOR,
+):
+    """
+    Registers the decorated function as the ``fixer`` of ``target_class_or_classes``, which is
+    the function that will be invoked every time you want to fix an incompatoble module to make
+    it work for training with Opacus.
+    You may supply your own validator_class that holds the registry of FIXERS.
+    The signature of every fixer is always the same:
+
+    >>> @register_module_fixer(nn.MyCustomClass)
+    >>> def fix(module: nn.Module, **kwargs) -> nn.Module:
+    >>>    pass
+
+    It may help you to take a look at the existing fixers inside Opacus, under ``opacus.validators.``
+    """
+
+    def decorator(f):
+        target_classes = (
+            target_class_or_classes
+            if isinstance(target_class_or_classes, Sequence)
+            else [target_class_or_classes]
+        )
+        for target_class in target_classes:
+            validator_class.FIXERS[target_class] = f
+        return f
+
+    return decorator


Test patch:
diff --git a/opacus/tests/dp_model_inspector_test.py b/opacus/tests/dp_model_inspector_test.py
deleted file mode 100644
index 5af59da7..00000000
--- a/opacus/tests/dp_model_inspector_test.py
+++ /dev/null
@@ -1,107 +0,0 @@
-#!/usr/bin/env python3
-# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
-
-
-import unittest
-
-import torch
-import torch.nn as nn
-from opacus import dp_model_inspector as dp_inspector
-from opacus.utils.module_modification import convert_batchnorm_modules
-from torchvision import models
-
-
-class dp_model_inspector_test(unittest.TestCase):
-    def test_raises_exception(self):
-        inspector = dp_inspector.DPModelInspector()
-        model = models.resnet50()
-        with self.assertRaises(dp_inspector.IncompatibleModuleException):
-            inspector.validate(model)
-
-    def test_returns_False(self):
-        inspector = dp_inspector.DPModelInspector(should_throw=False)
-        model = models.resnet50()
-        self.assertFalse(inspector.validate(model))
-
-    def test_raises_for_eval_mode(self):
-        inspector = dp_inspector.DPModelInspector()
-        model = models.resnet50()
-        model = model.eval()
-        with self.assertRaises(dp_inspector.IncompatibleModuleException):
-            inspector.validate(model)
-
-    def test_convert_batchnorm(self):
-        inspector = dp_inspector.DPModelInspector()
-        model = convert_batchnorm_modules(models.resnet50())
-        self.assertTrue(inspector.validate(model))
-
-    def test_running_stats(self):
-        inspector = dp_inspector.DPModelInspector(should_throw=False)
-
-        self.assertTrue(inspector.validate(nn.InstanceNorm1d(16)))
-        self.assertTrue(inspector.validate(nn.InstanceNorm1d(16, affine=True)))
-        self.assertTrue(
-            inspector.validate(nn.InstanceNorm1d(16, track_running_stats=True))
-        )
-        self.assertFalse(
-            inspector.validate(
-                nn.InstanceNorm1d(16, affine=True, track_running_stats=True)
-            )
-        )
-
-    def test_extra_param(self):
-        inspector = dp_inspector.DPModelInspector(should_throw=False)
-
-        class SampleNetWithExtraParam(nn.Module):
-            def __init__(self):
-                super().__init__()
-
-                self.fc = nn.Linear(8, 16)
-                self.extra_param = nn.Parameter(torch.Tensor(16, 2))
-
-            def forward(self, x):
-                x = self.fc(x)
-                x = x.matmul(self.extra_param)
-                return x
-
-        model = SampleNetWithExtraParam()
-        self.assertFalse(inspector.validate(model))
-
-        model.extra_param.requires_grad = False
-        self.assertTrue(inspector.validate(model))
-
-    def test_unsupported_layer(self):
-        class SampleNetWithTransformer(nn.Module):
-            def __init__(self):
-                super().__init__()
-
-                self.fc = nn.Linear(8, 16)
-                self.encoder = nn.Transformer()
-
-            def forward(self, x):
-                x = self.fc(x)
-                x = self.encoder(x)
-                return x
-
-        model = SampleNetWithTransformer()
-        inspector = dp_inspector.DPModelInspector(should_throw=False)
-        self.assertFalse(inspector.validate(model))
-
-    def test_conv2d(self):
-        inspector = dp_inspector.DPModelInspector(should_throw=False)
-
-        self.assertTrue(
-            inspector.validate(
-                nn.Conv2d(in_channels=3, out_channels=6, kernel_size=1, groups=1)
-            )
-        )
-        self.assertTrue(
-            inspector.validate(
-                nn.Conv2d(in_channels=3, out_channels=6, kernel_size=1, groups=3)
-            )
-        )
-        self.assertFalse(
-            inspector.validate(
-                nn.Conv2d(in_channels=6, out_channels=6, kernel_size=1, groups=2)
-            )
-        )
diff --git a/opacus/tests/module_validator_test.py b/opacus/tests/module_validator_test.py
new file mode 100644
index 00000000..a1f50940
--- /dev/null
+++ b/opacus/tests/module_validator_test.py
@@ -0,0 +1,68 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+import unittest
+
+import torch
+import torch.nn as nn
+from opacus.validators.errors import UnsupportedModuleError
+from opacus.validators.module_validator import ModuleValidator
+from torchvision.models import mobilenet_v3_small
+
+
+class ModuleValidator_test(unittest.TestCase):
+    def setUp(self):
+        self.original_model = mobilenet_v3_small()
+        self.fixed_model = ModuleValidator.fix(self.original_model)
+
+    def test_is_valid(self):
+        self.assertFalse(ModuleValidator.is_valid(self.original_model))
+        self.assertTrue(ModuleValidator.is_valid(self.fixed_model))
+
+    def test_validate_invalid_model(self):
+        with self.assertRaises(UnsupportedModuleError):
+            ModuleValidator.validate(self.original_model, raise_if_error=True)
+        errors = ModuleValidator.validate(self.original_model)
+        self.assertGreater(len(errors), 0)
+
+    def test_validate_valid_model(self):
+        errors = ModuleValidator.validate(self.fixed_model)
+        self.assertEqual(len(errors), 0)
+
+    def test_validate_training_mode(self):
+        self.fixed_model.eval()
+        self.assertFalse(ModuleValidator.is_valid(self.fixed_model))
+        self.fixed_model.train()
+        self.assertTrue(ModuleValidator.is_valid(self.fixed_model))
+
+    def test_is_valid_unsupported_grdsample_module(self):
+        unsupported_module = nn.Bilinear(2, 2, 2)  # currently not implemented
+        self.assertFalse(ModuleValidator.is_valid(unsupported_module))
+
+    def test_is_valid_extra_param(self):
+        class SampleNetWithExtraParam(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.fc = nn.Linear(8, 16)
+                self.extra_param = nn.Parameter(torch.Tensor(16, 2))
+
+            def forward(self, x):
+                x = self.fc(x)
+                x = x.matmul(self.extra_param)
+                return x
+
+        model = SampleNetWithExtraParam()
+        self.assertFalse(ModuleValidator.is_valid(model))
+
+        model.extra_param.requires_grad = False
+        self.assertTrue(ModuleValidator.is_valid(model))
+
+    def test_fix(self):
+        with self.assertLogs(level="INFO") as log_cm:
+            ModuleValidator.fix(self.original_model)
+            self.assertGreater(len(log_cm.records), 0)
+            for log_record in log_cm.records:
+                log_msg = log_record.getMessage()
+                self.assertRegex(log_msg, "Replaced sub_module .+ with .*")
+
+        # with self.assertNoLogs(level="INFO"):
+        #     ModuleValidator.fix(self.fixed_model)
diff --git a/opacus/tests/privacy_engine_test.py b/opacus/tests/privacy_engine_test.py
index 096c75e6..dc02f19c 100644
--- a/opacus/tests/privacy_engine_test.py
+++ b/opacus/tests/privacy_engine_test.py
@@ -9,7 +9,7 @@
 import torch.nn.functional as F
 from hypothesis import given, settings
 from opacus import PrivacyEngine
-from opacus.dp_model_inspector import IncompatibleModuleException
+from opacus.validators.errors import UnsupportedModuleError
 from torch.utils.data import DataLoader
 from torchvision import models, transforms
 from torchvision.datasets import FakeData
@@ -102,6 +102,7 @@ def _init_private_training(
         noise_multiplier: float = 1.0,
         max_grad_norm: float = 1.0,
         poisson_sampling: bool = True,
+        try_fix_incompatible_modules: bool = False,
     ):
         model = SampleConvNet()
         optimizer = torch.optim.SGD(model.parameters(), lr=self.LR, momentum=0)
@@ -119,6 +120,7 @@ def _init_private_training(
             noise_multiplier=noise_multiplier,
             max_grad_norm=max_grad_norm,
             poisson_sampling=poisson_sampling,
+            try_fix_incompatible_modules=try_fix_incompatible_modules,
         )
 
         return model, optimizer, dl, privacy_engine
@@ -282,25 +284,45 @@ def test_noise_changes_every_time(self):
         for p0, p1 in zip(first_run_params, second_run_params):
             self.assertFalse(torch.allclose(p0, p1))
 
-    @unittest.skip("Not yet implemented")
     def test_model_validator(self):
         """
-        Test that the privacy engine throws on attach
+        Test that the privacy engine raises errors
         if there are unsupported modules
         """
         resnet = models.resnet18()
         optimizer = torch.optim.SGD(resnet.parameters(), lr=1.0)
         privacy_engine = PrivacyEngine()
+        dl, _ = self._init_data()
 
-        with self.assertRaises(IncompatibleModuleException):
+        with self.assertRaises(UnsupportedModuleError):
             _, _, _ = privacy_engine.make_private(
                 module=resnet,
                 optimizer=optimizer,
-                data_loader=self.dl,
+                data_loader=dl,
                 noise_multiplier=1.3,
                 max_grad_norm=1,
             )
 
+    def test_model_validator_after_fix(self):
+        """
+        Test that the privacy engine fixes unsupported modules
+        and succeeds.
+        """
+        resnet = models.resnet18()
+        optimizer = torch.optim.SGD(resnet.parameters(), lr=1.0)
+        privacy_engine = PrivacyEngine()
+        dl, _ = self._init_data()
+
+        _, _, _ = privacy_engine.make_private(
+            module=resnet,
+            optimizer=optimizer,
+            data_loader=dl,
+            noise_multiplier=1.3,
+            max_grad_norm=1,
+            try_fix_incompatible_modules=True,
+        )
+        self.assertTrue(1, 1)
+
     def test_deterministic_run(self):
         """
         Tests that for 2 different models, secure seed can be fixed
@@ -332,7 +354,9 @@ def test_noise_level(self, noise_multiplier: float, max_steps: int):
         Tests that the noise level is correctly set
         """
         # Initialize models with parameters to zero
-        model, optimizer, dl, _ = self._init_private_training(noise_multiplier=noise_multiplier)
+        model, optimizer, dl, _ = self._init_private_training(
+            noise_multiplier=noise_multiplier
+        )
         for p in model.parameters():
             p.data.zero_()
 
@@ -353,12 +377,19 @@ def test_noise_level(self, noise_multiplier: float, max_steps: int):
                 break
 
         # Noise should be equal to lr*sigma*sqrt(n_params * steps) / batch_size
-        expected_norm = steps * n_params * optimizer.noise_multiplier**2 * self.LR**2 / (optimizer.expected_batch_size**2)
-        real_norm = sum([torch.sum(torch.pow(p.data, 2)) for p in model.parameters()]).item()
+        expected_norm = (
+            steps
+            * n_params
+            * optimizer.noise_multiplier ** 2
+            * self.LR ** 2
+            / (optimizer.expected_batch_size ** 2)
+        )
+        real_norm = sum(
+            [torch.sum(torch.pow(p.data, 2)) for p in model.parameters()]
+        ).item()
 
         self.assertAlmostEqual(real_norm, expected_norm, delta=0.05 * expected_norm)
 
-
     @unittest.skip("Not yet implemented")
     def test_raises_seed_set_on_secure_rng(self):
         """
diff --git a/opacus/tests/utils/__init__.py b/opacus/tests/utils/__init__.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/opacus/tests/utils/module_inspection_test.py b/opacus/tests/utils/module_inspection_test.py
deleted file mode 100644
index a19f624b..00000000
--- a/opacus/tests/utils/module_inspection_test.py
+++ /dev/null
@@ -1,110 +0,0 @@
-#!/usr/bin/env python3
-# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
-
-
-import unittest
-
-import torch.nn as nn
-from opacus.utils import module_inspection as mi
-from torchvision import models
-
-
-class utils_ModelInspector_test(unittest.TestCase):
-    def setUp(self):
-        def pred_supported(module):
-            return isinstance(module, (nn.Conv2d, nn.Linear))
-
-        def pred_not_unsupported(module):
-            return not isinstance(module, (nn.BatchNorm2d, nn.BatchNorm3d))
-
-        def pred_requires_grad(module):
-            return all(p.requires_grad for p in module.parameters(recurse=False))
-
-        self.pred_supported = pred_supported
-        self.pred_not_unsupported = pred_not_unsupported
-        self.pred_mix = lambda m: (not pred_requires_grad(m)) or pred_not_unsupported(m)
-
-    def test_validate_basic(self):
-        inspector = mi.ModelInspector(
-            "pred", lambda model: isinstance(model, nn.Linear)
-        )
-        model = nn.Conv1d(1, 1, 1)
-        valid = inspector.validate(model)
-        self.assertFalse(valid, inspector.violators)
-
-    def test_validate_positive_predicate_valid(self):
-        # test when a positive predicate (e.g. supported) returns true
-        inspector = mi.ModelInspector("pred", self.pred_supported)
-        model = nn.Conv2d(1, 1, 1)
-        valid = inspector.validate(model)
-        self.assertTrue(valid)
-        list_len = len(inspector.violators)
-        self.assertEqual(list_len, 0, f"violators = {inspector.violators}")
-
-    def test_validate_positive_predicate_invalid(self):
-        # test when a positive predicate (e.g. supported) returns false
-        inspector = mi.ModelInspector("pred", self.pred_supported)
-        model = nn.Conv1d(1, 1, 1)
-        valid = inspector.validate(model)
-        self.assertFalse(valid)
-        list_len = len(inspector.violators)
-        self.assertEqual(list_len, 1, f"violators = {inspector.violators}")
-
-    def test_validate_negative_predicate_ture(self):
-        # test when a negative predicate (e.g. not unsupported) returns true
-        inspector = mi.ModelInspector("pred1", self.pred_not_unsupported)
-        model = nn.Sequential(nn.Conv2d(1, 1, 1), nn.Linear(1, 1))
-        valid = inspector.validate(model)
-        self.assertTrue(valid)
-        list_len = len(inspector.violators)
-        self.assertEqual(list_len, 0)
-
-    def test_validate_negative_predicate_False(self):
-        # test when a negative predicate (e.g. not unsupported) returns false
-        inspector = mi.ModelInspector("pred", self.pred_not_unsupported)
-        model = nn.Sequential(nn.Conv2d(1, 1, 1), nn.BatchNorm2d(1))
-        valid = inspector.validate(model)
-        self.assertFalse(valid)
-        list_len = len(inspector.violators)
-        self.assertEqual(list_len, 1, f"violators = {inspector.violators}")
-
-    def test_validate_mix_predicate(self):
-        # check with a mix predicate not requires grad or is not unsupported
-        inspector = mi.ModelInspector("pred1", self.pred_mix)
-        model = nn.Sequential(nn.Conv2d(1, 1, 1), nn.BatchNorm2d(1))
-        for p in model[1].parameters():
-            p.requires_grad = False
-        valid = inspector.validate(model)
-        self.assertTrue(valid)
-
-    def test_check_everything_flag(self):
-        # check to see if a model does not containt nn.sequential
-        inspector = mi.ModelInspector(
-            "pred",
-            lambda model: not isinstance(model, nn.Sequential),
-            check_leaf_nodes_only=False,
-        )
-        model = nn.Sequential(nn.Conv1d(1, 1, 1))
-        valid = inspector.validate(model)
-        self.assertFalse(valid, f"violators = {inspector.violators}")
-
-    def test_complicated_case(self):
-        def good(x):
-            return isinstance(x, (nn.Conv2d, nn.Linear))
-
-        def bad(x):
-            return isinstance(x, nn.modules.batchnorm._BatchNorm)
-
-        inspector1 = mi.ModelInspector("good_or_bad", lambda x: good(x) or bad(x))
-        inspector2 = mi.ModelInspector("not_bad", lambda x: not bad(x))
-        model = models.resnet50()
-        valid = inspector1.validate(model)
-        self.assertTrue(valid, f"violators = {inspector1.violators}")
-        self.assertEqual(
-            len(inspector1.violators), 0, f"violators = {inspector1.violators}"
-        )
-        valid = inspector2.validate(model)
-        self.assertFalse(valid, f"violators = {inspector2.violators}")
-        self.assertEqual(
-            len(inspector2.violators), 53, f"violators = {inspector2.violators}"
-        )
diff --git a/opacus/tests/utils/module_modification_test.py b/opacus/tests/utils/module_modification_test.py
deleted file mode 100644
index 49e6dc6c..00000000
--- a/opacus/tests/utils/module_modification_test.py
+++ /dev/null
@@ -1,176 +0,0 @@
-#!/usr/bin/env python3
-# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
-
-
-import unittest
-from typing import Tuple
-
-import torch
-import torch.nn as nn
-from opacus import PrivacyEngine
-from opacus.dp_model_inspector import IncompatibleModuleException
-from opacus.utils import module_modification as mm
-from torch.utils.data import DataLoader
-from torchvision import models, transforms
-from torchvision.datasets import FakeData
-
-
-class ReplaceAllModulesTest(unittest.TestCase):
-    def checkModulePresent(self, root: nn.Module, targetclass):
-        result = any(isinstance(module, targetclass) for module in root.modules())
-        self.assertTrue(result)
-
-    def checkModuleNotPresent(self, root: nn.Module, targetclass):
-        for module in root.modules():
-            self.assertFalse(
-                isinstance(module, targetclass),
-                msg=f"{module} has the given targetclass type",
-            )
-
-    def test_replace_basic_case(self):
-        model = nn.BatchNorm1d(10)
-        model = mm.replace_all_modules(
-            model, nn.BatchNorm1d, lambda _: nn.BatchNorm2d(10)
-        )
-        self.checkModulePresent(model, nn.BatchNorm2d)
-        self.checkModuleNotPresent(model, nn.BatchNorm1d)
-
-    def test_replace_sequential_case(self):
-        model = nn.Sequential(nn.Conv1d(1, 2, 3), nn.Sequential(nn.Conv2d(4, 5, 6)))
-
-        def conv(m: nn.Conv2d):
-            return nn.Linear(4, 5)
-
-        model = mm.replace_all_modules(model, nn.Conv2d, conv)
-        self.checkModulePresent(model, nn.Linear)
-        self.checkModuleNotPresent(model, nn.Conv2d)
-
-    def test_nullify_resnet18(self):
-        model = models.resnet18()
-        # check module BatchNorms is there
-        self.checkModulePresent(model, nn.BatchNorm2d)
-        # nullify the module (replace with Idetity)
-        model = mm.nullify_batchnorm_modules(model)
-        # check module is not present
-        self.checkModuleNotPresent(model, nn.BatchNorm2d)
-
-    def test_convert_batchnorm_modules_resnet50(self):
-        model = models.resnet50()
-        # check module BatchNorms is there
-        self.checkModulePresent(model, nn.BatchNorm2d)
-        # replace the module with instancenorm
-        model = mm.convert_batchnorm_modules(model)
-        # check module is not present
-        self.checkModuleNotPresent(model, nn.BatchNorm2d)
-        self.checkModulePresent(model, nn.GroupNorm)
-
-
-class BasicModel(nn.Module):
-    def __init__(self, imgSize):
-        super().__init__()
-        self.size = imgSize[0] * imgSize[1] * imgSize[2]
-        self.bn = nn.BatchNorm2d(imgSize[0])
-        self.fc = nn.Linear(self.size, 2)
-
-    def forward(self, input):
-        x = self.bn(input)
-        x = x.view(-1, self.size)
-        x = self.fc(x)
-        return x
-
-
-class ConvertBatchnormModulesTest(unittest.TestCase):
-    def setUp(self):
-        self.criterion = nn.CrossEntropyLoss()
-
-    def genFakeData(
-        self, imgSize: Tuple[int, int, int], batch_size: int = 1, num_batches: int = 1
-    ) -> DataLoader:
-        self.ds = FakeData(
-            size=num_batches,
-            image_size=imgSize,
-            num_classes=2,
-            transform=transforms.Compose([transforms.ToTensor()]),
-        )
-        return DataLoader(self.ds, batch_size=batch_size)
-
-    def runOneBatch(
-        self,
-        model: nn.Module,
-        imgsize: Tuple[int, int, int],
-    ):
-        dl = self.genFakeData(imgsize, 1, 1)
-        optimizer = torch.optim.SGD(model.parameters(), lr=1e-2, momentum=0.9)
-
-        privacy_engine = PrivacyEngine()
-        model, optimizer, dl = privacy_engine.make_private(
-            module=model,
-            optimizer=optimizer,
-            data_loader=dl,
-            noise_multiplier=1.0,
-            max_grad_norm=1.0,
-        )
-
-        for x, y in dl:
-            # forward
-            try:
-                logits = model(x)
-            except Exception as err:
-                self.fail(f"Failed forward step with exception: {err}")
-            loss = self.criterion(logits, y)
-            # backward
-            try:
-                loss.backward()
-            except Exception as err:
-                self.fail(f"Failed backward step with exception: {err}")
-            # optimizer
-            try:
-                optimizer.step()
-            except Exception as err:
-                self.fail(f"Failed optimizer step with exception: {err}")
-        optimizer.zero_grad()
-
-    @unittest.skip("Not yet implemented")
-    def test_run_basic_case(self):
-        imgSize = (3, 4, 5)
-        # should throw because privacy engine does not work with batch norm
-        # remove the next two lines when we support batch norm
-        with self.assertRaises(IncompatibleModuleException):
-            self.runOneBatch(BasicModel(imgSize), imgSize)
-        self.runOneBatch(mm.convert_batchnorm_modules(BasicModel(imgSize)), imgSize)
-
-    @unittest.skip("Not yet implemented")
-    def test_run_resnet18(self):
-        imgSize = (3, 224, 224)
-        # should throw because privacy engine does not work with batch norm
-        # remove the next two lines when we support batch norm
-        with self.assertRaises(IncompatibleModuleException):
-            self.runOneBatch(models.resnet18(), imgSize)
-        self.runOneBatch(mm.convert_batchnorm_modules(models.resnet18()), imgSize)
-
-    @unittest.skip("Not yet implemented")
-    def test_run_resnet34(self):
-        imgSize = (3, 224, 224)
-        # should throw because privacy engine does not work with batch norm
-        # remove the next two lines when we support batch norm
-        with self.assertRaises(IncompatibleModuleException):
-            self.runOneBatch(models.resnet34(), imgSize)
-        self.runOneBatch(mm.convert_batchnorm_modules(models.resnet34()), imgSize)
-
-    @unittest.skip("Not yet implemented")
-    def test_run_resnet50(self):
-        imgSize = (3, 224, 224)
-        # should throw because privacy engine does not work with batch norm
-        # remove the next two lines when we support batch norm
-        with self.assertRaises(IncompatibleModuleException):
-            self.runOneBatch(models.resnet50(), imgSize)
-        self.runOneBatch(mm.convert_batchnorm_modules(models.resnet50()), imgSize)
-
-    @unittest.skip("Not yet implemented")
-    def test_run_resnet101(self):
-        imgSize = (3, 224, 224)
-        # should throw because privacy engine does not work with batch norm
-        # remove the next two lines when we support batch norm
-        with self.assertRaises(IncompatibleModuleException):
-            self.runOneBatch(models.resnet101(), imgSize)
-        self.runOneBatch(mm.convert_batchnorm_modules(models.resnet101()), imgSize)
diff --git a/opacus/tests/validators/__init__.py b/opacus/tests/validators/__init__.py
new file mode 100644
index 00000000..4e3dc62d
--- /dev/null
+++ b/opacus/tests/validators/__init__.py
@@ -0,0 +1,2 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
diff --git a/opacus/tests/validators/batch_norm_test.py b/opacus/tests/validators/batch_norm_test.py
new file mode 100644
index 00000000..4dfdc02c
--- /dev/null
+++ b/opacus/tests/validators/batch_norm_test.py
@@ -0,0 +1,44 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+import unittest
+
+import torch.nn as nn
+from opacus.validators.errors import ShouldReplaceModuleError
+from opacus.validators.module_validator import ModuleValidator
+
+
+class BatchNormValidator_test(unittest.TestCase):
+    def setUp(self):
+        self.bn1 = nn.BatchNorm1d(4)
+        self.bn2 = nn.BatchNorm2d(4)
+        self.bn3 = nn.BatchNorm3d(4)
+        self.bns = nn.SyncBatchNorm(4)
+        self.mv = ModuleValidator.VALIDATORS
+        self.mf = ModuleValidator.FIXERS
+
+    def test_validate(self):
+        val1 = self.mv[type(self.bn1)](self.bn1)
+        val2 = self.mv[type(self.bn2)](self.bn2)
+        val3 = self.mv[type(self.bn3)](self.bn3)
+        vals = self.mv[type(self.bns)](self.bns)
+
+        self.assertEqual(len(val1), 1)
+        self.assertEqual(len(val2), 1)
+        self.assertEqual(len(val3), 1)
+        self.assertEqual(len(vals), 1)
+
+        self.assertTrue(isinstance(val1[0], ShouldReplaceModuleError))
+        self.assertTrue(isinstance(val2[0], ShouldReplaceModuleError))
+        self.assertTrue(isinstance(val3[0], ShouldReplaceModuleError))
+        self.assertTrue(isinstance(vals[0], ShouldReplaceModuleError))
+
+    def test_fix(self):
+        fix1 = self.mf[type(self.bn1)](self.bn1)
+        fix2 = self.mf[type(self.bn2)](self.bn2)
+        fix3 = self.mf[type(self.bn3)](self.bn3)
+        fixs = self.mf[type(self.bns)](self.bns)
+
+        self.assertTrue(isinstance(fix1, nn.GroupNorm))
+        self.assertTrue(isinstance(fix2, nn.GroupNorm))
+        self.assertTrue(isinstance(fix3, nn.GroupNorm))
+        self.assertTrue(isinstance(fixs, nn.GroupNorm))
diff --git a/opacus/tests/validators/ddp_test.py b/opacus/tests/validators/ddp_test.py
new file mode 100644
index 00000000..60cf0ff6
--- /dev/null
+++ b/opacus/tests/validators/ddp_test.py
@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+import os
+import unittest
+
+import torch.distributed as dist
+import torch.nn as nn
+from opacus.validators.errors import ShouldReplaceModuleError
+from opacus.validators.errors import ShouldReplaceModuleError
+from opacus.validators.module_validator import ModuleValidator
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+
+class DDPValidator_test(unittest.TestCase):
+    def setUp(self):
+        self.module = nn.Linear(8, 4)
+        self.mv = ModuleValidator.VALIDATORS
+        self.mf = ModuleValidator.FIXERS
+        self._setup_dist()
+
+    def tearDown(self):
+        self._cleanup_dist()
+
+    def _cleanup_dist(self):
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+    def _setup_dist(self):
+        os.environ["MASTER_ADDR"] = "localhost"
+        os.environ["MASTER_PORT"] = "12355"
+        # initialize the process group
+        dist.init_process_group("gloo", rank=0, world_size=1)
+
+    def test_validate(self):
+        ddp = DDP(self.module)
+        val_ddp = self.mv[type(ddp)](ddp)
+        self.assertTrue(len(val_ddp), 1)
+        self.assertTrue(isinstance(val_ddp[0], ShouldReplaceModuleError))
+
+    def test_fix(self):
+        ddp = DDP(self.module)
+        with self.assertRaises(ShouldReplaceModuleError):
+            fix_ddp = self.mf[type(ddp)](ddp)
diff --git a/opacus/tests/validators/instance_norm_test.py b/opacus/tests/validators/instance_norm_test.py
new file mode 100644
index 00000000..d64b5a60
--- /dev/null
+++ b/opacus/tests/validators/instance_norm_test.py
@@ -0,0 +1,51 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+import unittest
+
+import torch.nn as nn
+from opacus.validators.errors import IllegalModuleConfigurationError
+from opacus.validators.module_validator import ModuleValidator
+
+
+class InstanceNormValidator_test(unittest.TestCase):
+    def setUp(self):
+        self.in1 = nn.InstanceNorm1d(4, affine=True, track_running_stats=True)
+        self.in2 = nn.InstanceNorm2d(4, affine=False, track_running_stats=True)
+        self.in3 = nn.InstanceNorm3d(4, affine=False, track_running_stats=True)
+        self.in_no_stats = nn.InstanceNorm3d(4, affine=True)
+
+        self.mv = ModuleValidator.VALIDATORS
+        self.mf = ModuleValidator.FIXERS
+
+    def test_validate(self):
+        val1 = self.mv[type(self.in1)](self.in1)
+        val2 = self.mv[type(self.in2)](self.in2)
+        val3 = self.mv[type(self.in3)](self.in3)
+        vals = self.mv[type(self.in_no_stats)](self.in_no_stats)
+
+        self.assertEqual(len(val1), 1)
+        self.assertEqual(len(val2), 1)
+        self.assertEqual(len(val3), 1)
+        self.assertEqual(len(vals), 0)
+
+        self.assertTrue(isinstance(val1[0], IllegalModuleConfigurationError))
+        self.assertTrue(isinstance(val2[0], IllegalModuleConfigurationError))
+        self.assertTrue(isinstance(val3[0], IllegalModuleConfigurationError))
+
+    def test_fix(self):
+        fix1 = self.mf[type(self.in1)](self.in1)
+        fix2 = self.mf[type(self.in2)](self.in2)
+        fix3 = self.mf[type(self.in3)](self.in3)
+        fixs = self.mf[type(self.in_no_stats)](self.in_no_stats)
+
+        self.assertTrue(isinstance(fix1, nn.InstanceNorm1d))
+        self.assertFalse(fix1.track_running_stats)
+
+        self.assertTrue(isinstance(fix2, nn.InstanceNorm2d))
+        self.assertFalse(fix2.track_running_stats)
+
+        self.assertTrue(isinstance(fix3, nn.InstanceNorm3d))
+        self.assertFalse(fix3.track_running_stats)
+
+        self.assertTrue(isinstance(fixs, nn.InstanceNorm3d))
+        self.assertTrue(fixs is self.in_no_stats)
diff --git a/opacus/tests/validators/lstm_test.py b/opacus/tests/validators/lstm_test.py
new file mode 100644
index 00000000..dfdcfe6e
--- /dev/null
+++ b/opacus/tests/validators/lstm_test.py
@@ -0,0 +1,28 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+import unittest
+
+import torch.nn as nn
+from opacus.layers import DPLSTM
+from opacus.utils.module_utils import are_state_dict_equal
+from opacus.validators.errors import ShouldReplaceModuleError
+from opacus.validators.module_validator import ModuleValidator
+
+
+class LSTMValidator_test(unittest.TestCase):
+    def setUp(self):
+        self.lstm = nn.LSTM(8, 4)
+        self.mv = ModuleValidator.VALIDATORS
+        self.mf = ModuleValidator.FIXERS
+
+    def test_validate(self):
+        val_lstm = self.mv[type(self.lstm)](self.lstm)
+        self.assertEqual(len(val_lstm), 1)
+        self.assertTrue(isinstance(val_lstm[0], ShouldReplaceModuleError))
+
+    def test_fix(self):
+        fix_lstm = self.mf[type(self.lstm)](self.lstm)
+        self.assertTrue(isinstance(fix_lstm, DPLSTM))
+        self.assertTrue(
+            are_state_dict_equal(self.lstm.state_dict(), fix_lstm.state_dict())
+        )
diff --git a/opacus/tests/validators/multihead_attention_test.py b/opacus/tests/validators/multihead_attention_test.py
new file mode 100644
index 00000000..e97f2615
--- /dev/null
+++ b/opacus/tests/validators/multihead_attention_test.py
@@ -0,0 +1,28 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+import unittest
+
+import torch.nn as nn
+from opacus.layers import DPMultiheadAttention
+from opacus.utils.module_utils import are_state_dict_equal
+from opacus.validators.errors import ShouldReplaceModuleError
+from opacus.validators.module_validator import ModuleValidator
+
+
+class MultiheadAttentionValidator_test(unittest.TestCase):
+    def setUp(self):
+        self.mha = nn.MultiheadAttention(8, 4)
+        self.mv = ModuleValidator.VALIDATORS
+        self.mf = ModuleValidator.FIXERS
+
+    def test_validate(self):
+        val_mha = self.mv[type(self.mha)](self.mha)
+        self.assertEqual(len(val_mha), 1)
+        self.assertTrue(isinstance(val_mha[0], ShouldReplaceModuleError))
+
+    def test_fix(self):
+        fix_mha = self.mf[type(self.mha)](self.mha)
+        self.assertTrue(isinstance(fix_mha, DPMultiheadAttention))
+        self.assertFalse(  # state_dicts are not same
+            are_state_dict_equal(self.mha.state_dict(), fix_mha.state_dict())
+        )