Repository: TorchEnsemble-Community/Ensemble-Pytorch (language: python)

Issue:
Update the experimental branch



Discussion / hints:

Fix patch:
diff --git a/docs/introduction.rst b/docs/introduction.rst
index 0bfd77b..f2d1ba7 100644
--- a/docs/introduction.rst
+++ b/docs/introduction.rst
@@ -59,6 +59,16 @@ To obtain snapshots with good performance, snapshot ensemble uses **cyclic annea
 
 Notice that the iteration above indicates the loop on enumerating all batches within each epoch, instead of the loop on iterating over all training epochs.
 
+Adversarial Training [2]_
+-------------------------
+
+Adversarial samples can be used to improve the performance of base estimators, as validated by the authors in [2]. The implemented ``AdversarialTrainingClassifier`` and ``AdversarialTrainingRegressor`` contain :math:`M` independent base estimators, and each of them is fitted independently as in Voting and Bagging.
+
+During the training stage of each base estimator :math:`h^m`, an adversarial sample :math:`\mathbf{x}_i^{adv}` is first generated for each sample :math:`\mathbf{x}_i` in the current data batch, using the fast gradient sign method (FGSM). After then, parameters of the base estimator is optimized to minimize the training loss :math:`\mathcal{L}(\mathbf{o}_i, y_i) + \mathcal{L}(\mathbf{o}_i^{adv}, y_i)`, where :math:`\mathbf{o}_i^{adv}` is the output of :math:`h^m` on the adversarial sample :math:`\mathbf{x}_i^{adv}`. Clearly, this training loss encourages each base estimator to perform well on both original samples and adversarial samples.
+
+Same as Voting and Bagging, the output of ``AdversarialTrainingClassifier`` or ``AdversarialTrainingRegressor`` during the evaluating stage is the average over predictions from all base estimators.
+
 **References**
 
-.. [1] Huang, Gao, et al. "Snapshot ensembles: Train 1, get m for free." ICLR, 2017.
\ No newline at end of file
+.. [1] Huang Gao, Sharon Yixuan Li, Geoff Pleisset, et al., "Snapshot ensembles: Train 1, get m for free." ICLR, 2017.
+.. [2] Balaji Lakshminarayanan, Alexander Pritzel, Charles Blundell., "Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles." NIPS 2017.
\ No newline at end of file
diff --git a/docs/parameters.rst b/docs/parameters.rst
index 6133a1a..8fbbe89 100644
--- a/docs/parameters.rst
+++ b/docs/parameters.rst
@@ -1,7 +1,7 @@
 Parameters
 ==========
 
-This page gives the API reference of ``torchensemble``, please also refer to `Introduction <./introduction.html>`__ for details.
+This page provides the API reference of ``torchensemble``, please also refer to `Introduction <./introduction.html>`__ for details.
 
 Fusion
 ------
@@ -47,7 +47,7 @@ Bagging
 
 In bagging-based ensemble methods, each base estimator is trained
 independently. In addition, sampling with replacement is conducted on the
-training data to further encourge the diversity between different base
+training data to further encourage the diversity between different base
 estimators in the ensemble model.
 
 BaggingClassifier
@@ -108,3 +108,26 @@ SnapshotEnsembleRegressor
 
 .. autoclass:: torchensemble.snapshot_ensemble.SnapshotEnsembleRegressor
     :members:
+
+Adversarial Training
+--------------------
+
+Adversarial training is able to improve the performance of an ensemble by
+treating adversarial samples as the augmented training data. The fast
+gradient sign method (FGSM) is used to generate adversarial samples.
+
+Reference:
+    B. Lakshminarayanan, A. Pritzel, C. Blundell., Simple and Scalable
+    Predictive Uncertainty Estimation using Deep Ensembles, NIPS 2017.
+
+AdversarialTrainingClassifier
+*****************************
+
+.. autoclass:: torchensemble.adversarial_training.AdversarialTrainingClassifier
+    :members:
+
+AdversarialTrainingRegressor
+*****************************
+
+.. autoclass:: torchensemble.adversarial_training.AdversarialTrainingRegressor
+    :members:
diff --git a/torchensemble/__init__.py b/torchensemble/__init__.py
index 99392d0..85c4611 100644
--- a/torchensemble/__init__.py
+++ b/torchensemble/__init__.py
@@ -8,6 +8,8 @@
 from .gradient_boosting import GradientBoostingRegressor
 from .snapshot_ensemble import SnapshotEnsembleClassifier
 from .snapshot_ensemble import SnapshotEnsembleRegressor
+from .adversarial_training import AdversarialTrainingClassifier
+from .adversarial_training import AdversarialTrainingRegressor
 
 
 __all__ = ["FusionClassifier",
@@ -19,4 +21,6 @@
            "GradientBoostingClassifier",
            "GradientBoostingRegressor",
            "SnapshotEnsembleClassifier",
-           "SnapshotEnsembleRegressor"]
+           "SnapshotEnsembleRegressor",
+           "AdversarialTrainingClassifier",
+           "AdversarialTrainingRegressor"]
diff --git a/torchensemble/adversarial_training.py b/torchensemble/adversarial_training.py
new file mode 100644
index 0000000..4b5f82c
--- /dev/null
+++ b/torchensemble/adversarial_training.py
@@ -0,0 +1,470 @@
+"""
+  Adversarial training is able to improve the performance of an ensemble by
+  treating adversarial samples as the augmented training data. The fast
+  gradient sign method (FGSM) is used to generate adversarial samples.
+
+  Reference:
+      B. Lakshminarayanan, A. Pritzel, C. Blundell., Simple and Scalable
+      Predictive Uncertainty Estimation using Deep Ensembles, NIPS 2017.
+"""
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from joblib import Parallel, delayed
+
+from ._base import BaseModule, torchensemble_model_doc
+from . import utils
+
+
+__all__ = ["_BaseAdversarialTraining",
+           "AdversarialTrainingClassifier",
+           "AdversarialTrainingRegressor"]
+
+
+__fit_doc = """
+    Parameters
+    ----------
+    train_loader : torch.utils.data.DataLoader
+        A :mod:`DataLoader` container that contains the training data.
+    lr : float, default=1e-3
+        The learning rate of the parameter optimizer.
+    weight_decay : float, default=5e-4
+        The weight decay of the parameter optimizer.
+    epochs : int, default=100
+        The number of training epochs.
+    optimizer : {"SGD", "Adam", "RMSprop"}, default="Adam"
+        The type of parameter optimizer.
+    epsilon : float, defaul=0.01
+        The step used to generate adversarial samples in the fast gradient
+        sign method (FGSM), which should be in the range [0, 1].
+    log_interval : int, default=100
+        The number of batches to wait before printting the training status.
+    test_loader : torch.utils.data.DataLoader, default=None
+        A :mod:`DataLoader` container that contains the evaluating data.
+
+        - If ``None``, no validation is conducted after each training
+          epoch.
+        - If not ``None``, the ensemble will be evaluated on this
+          dataloader after each training epoch.
+    save_model : bool, default=True
+        Whether to save the model.
+
+        - If test_loader is ``None``, the ensemble containing
+          ``n_estimators`` base estimators will be saved.
+        - If test_loader is not ``None``, the ensemble with the best
+          validation performance will be saved.
+    save_dir : string, default=None
+        Specify where to save the model.
+
+        - If ``None``, the model will be saved in the current directory.
+        - If not ``None``, the model will be saved in the specified
+          directory: ``save_dir``.
+"""
+
+
+def _adversarial_training_model_doc(header, item="fit"):
+    """
+    Decorator on obtaining documentation for different adversarial training
+    models.
+    """
+    def get_doc(item):
+        """Return selected item"""
+        __doc = {"fit": __fit_doc}
+        return __doc[item]
+
+    def adddoc(cls):
+        doc = [header + "\n\n"]
+        doc.extend(get_doc(item))
+        cls.__doc__ = "".join(doc)
+        return cls
+    return adddoc
+
+
+def _parallel_fit_per_epoch(train_loader,
+                            lr,
+                            weight_decay,
+                            epoch,
+                            optimizer,
+                            epsilon,
+                            log_interval,
+                            idx,
+                            estimator,
+                            criterion,
+                            device,
+                            is_classification):
+    """Private function used to fit base estimators in parallel."""
+    optimizer = utils.set_optimizer(estimator, optimizer, lr, weight_decay)
+
+    for batch_idx, (data, target) in enumerate(train_loader):
+
+        batch_size = data.size()[0]
+        data, target = data.to(device), target.to(device)
+        data.requires_grad = True
+
+        # Get adversarial samples
+        _output = estimator(data)
+        _loss = criterion(_output, target)
+        _loss.backward()
+        data_grad = data.grad.data
+        adv_data = _get_fgsm_samples(data, epsilon, data_grad)
+
+        # Compute the training loss
+        optimizer.zero_grad()
+        org_output = estimator(data)
+        adv_output = estimator(adv_data)
+        loss = criterion(org_output, target) + criterion(adv_output, target)
+        loss.backward()
+        optimizer.step()
+
+        # Print training status
+        if batch_idx % log_interval == 0:
+
+            # Classification
+            if is_classification:
+                _, predicted = torch.max(org_output.data, 1)
+                correct = (predicted == target).sum().item()
+
+                msg = ("Estimator: {:03d} | Epoch: {:03d} | Batch: {:03d}"
+                       " | Loss: {:.5f} | Correct: {:d}/{:d}")
+                print(
+                    msg.format(
+                        idx, epoch, batch_idx, loss, correct, batch_size
+                    )
+                )
+            # Regression
+            else:
+                msg = ("Estimator: {:03d} | Epoch: {:03d} | Batch: {:03d}"
+                       " | Loss: {:.5f}")
+                print(msg.format(idx, epoch, batch_idx, loss))
+
+    return estimator
+
+
+def _get_fgsm_samples(sample, epsilon, sample_grad):
+    """
+    Private functions used to generate adversarial samples with fast gradient
+    sign method (FGSM)."""
+
+    # Check the input range of `sample`
+    min_value, max_value = torch.min(sample), torch.max(sample)
+    if not 0 <= min_value < max_value <= 1:
+        msg = ("The input range of samples passed to adversarial training"
+               " should be in the range [0, 1], but got [{:.3f}, {:.3f}]"
+               " instead.")
+        raise ValueError(msg.format(min_value, max_value))
+
+    sign_sample_grad = sample_grad.sign()
+    perturbed_sample = sample + epsilon * sign_sample_grad
+    perturbed_sample = torch.clamp(perturbed_sample, 0, 1)
+    return perturbed_sample
+
+
+class _BaseAdversarialTraining(BaseModule):
+
+    def _validate_parameters(self,
+                             lr,
+                             weight_decay,
+                             epochs,
+                             epsilon,
+                             log_interval):
+        """Validate hyper-parameters on training the ensemble."""
+
+        if not lr > 0:
+            msg = ("The learning rate of optimizer = {} should be strictly"
+                   " positive.")
+            self.logger.error(msg.format(lr))
+            raise ValueError(msg.format(lr))
+
+        if not weight_decay >= 0:
+            msg = "The weight decay of optimizer = {} should not be negative."
+            self.logger.error(msg.format(weight_decay))
+            raise ValueError(msg.format(weight_decay))
+
+        if not epochs > 0:
+            msg = ("The number of training epochs = {} should be strictly"
+                   " positive.")
+            self.logger.error(msg.format(epochs))
+            raise ValueError(msg.format(epochs))
+
+        if not 0 < epsilon <= 1:
+            msg = ("The step used to generate adversarial samples in FGSM"
+                   " should be in the range (0, 1], but got {} instead.")
+            self.logger.error(msg.format(epsilon))
+            raise ValueError(msg.format(epsilon))
+
+        if not log_interval > 0:
+            msg = ("The number of batches to wait before printting the"
+                   " training status should be strictly positive, but got {}"
+                   " instead.")
+            self.logger.error(msg.format(log_interval))
+            raise ValueError(msg.format(log_interval))
+
+
+@torchensemble_model_doc("""Implementation on the AdversarialTrainingClassifier.""",  # noqa: E501
+                         "model")
+class AdversarialTrainingClassifier(_BaseAdversarialTraining):
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.is_classification = True
+
+    @torchensemble_model_doc(
+        """Implementation on the data forwarding in AdversarialTrainingClassifier.""",  # noqa: E501
+        "classifier_forward")
+    def forward(self, x):
+        batch_size = x.size(0)
+        proba = torch.zeros(batch_size, self.n_outputs).to(self.device)
+
+        # Take the average over class distributions from all base estimators.
+        for estimator in self.estimators_:
+            proba += F.softmax(estimator(x), dim=1) / self.n_estimators
+
+        return proba
+
+    @_adversarial_training_model_doc(
+        """Implementation on the training stage of AdversarialTrainingClassifier.""",  # noqa: E501
+        "fit"
+    )
+    def fit(self,
+            train_loader,
+            lr=1e-3,
+            weight_decay=5e-4,
+            epochs=100,
+            optimizer="Adam",
+            epsilon=0.5,
+            log_interval=100,
+            test_loader=None,
+            save_model=True,
+            save_dir=None):
+
+        # Instantiate base estimators and set attributes
+        estimators = []
+        for _ in range(self.n_estimators):
+            estimators.append(self._make_estimator())
+        self._validate_parameters(lr,
+                                  weight_decay,
+                                  epochs,
+                                  epsilon,
+                                  log_interval)
+        self.n_outputs = self._decide_n_outputs(train_loader, True)
+
+        # Utils
+        criterion = nn.CrossEntropyLoss()
+        best_acc = 0.
+
+        # Internal helper function on pesudo forward
+        def _forward(estimators, data):
+            batch_size = data.size()[0]
+            proba = torch.zeros(batch_size, self.n_outputs).to(self.device)
+
+            for estimator in estimators:
+                proba += F.softmax(estimator(data), dim=1) / self.n_estimators
+
+            return proba
+
+        # Maintain a pool of workers
+        with Parallel(n_jobs=self.n_jobs) as parallel:
+
+            # Training loop
+            for epoch in range(epochs):
+                self.train()
+                rets = parallel(delayed(_parallel_fit_per_epoch)(
+                        train_loader,
+                        lr,
+                        weight_decay,
+                        epoch,
+                        optimizer,
+                        epsilon,
+                        log_interval,
+                        idx,
+                        estimator,
+                        criterion,
+                        self.device,
+                        True
+                    )
+                    for idx, estimator in enumerate(estimators)
+                )
+
+                estimators = rets  # update
+
+                # Validation
+                if test_loader:
+                    self.eval()
+                    with torch.no_grad():
+                        correct = 0
+                        total = 0
+                        for _, (data, target) in enumerate(test_loader):
+                            data = data.to(self.device)
+                            target = target.to(self.device)
+                            output = _forward(estimators, data)
+                            _, predicted = torch.max(output.data, 1)
+                            correct += (predicted == target).sum().item()
+                            total += target.size(0)
+                        acc = 100 * correct / total
+
+                        if acc > best_acc:
+                            best_acc = acc
+                            self.estimators_ = nn.ModuleList()  # reset
+                            self.estimators_.extend(estimators)
+                            if save_model:
+                                utils.save(self, save_dir, self.logger)
+
+                        msg = ("Epoch: {:03d} | Validation Acc: {:.3f}"
+                               " % | Historical Best: {:.3f} %")
+                        self.logger.info(msg.format(epoch, acc, best_acc))
+
+        self.estimators_ = nn.ModuleList()
+        self.estimators_.extend(rets)
+        if save_model and not test_loader:
+            utils.save(self, save_dir, self.logger)
+
+    @torchensemble_model_doc(
+        """Implementation on the evaluating stage of AdversarialTrainingClassifier.""",  # noqa: E501
+        "classifier_predict")
+    def predict(self, test_loader):
+        self.eval()
+        correct = 0
+        total = 0
+
+        for _, (data, target) in enumerate(test_loader):
+            data, target = data.to(self.device), target.to(self.device)
+            output = self.forward(data)
+            _, predicted = torch.max(output.data, 1)
+            correct += (predicted == target).sum().item()
+            total += target.size(0)
+
+        acc = 100 * correct / total
+
+        return acc
+
+
+@torchensemble_model_doc("""Implementation on the AdversarialTrainingRegressor.""",  # noqa: E501
+                         "model")
+class AdversarialTrainingRegressor(_BaseAdversarialTraining):
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.is_classification = False
+
+    @torchensemble_model_doc(
+        """Implementation on the data forwarding in AdversarialTrainingRegressor.""",  # noqa: E501
+        "regressor_forward")
+    def forward(self, x):
+        batch_size = x.size(0)
+        pred = torch.zeros(batch_size, self.n_outputs).to(self.device)
+
+        # Take the average over predictions from all base estimators.
+        for estimator in self.estimators_:
+            pred += estimator(x) / self.n_estimators
+
+        return pred
+
+    @_adversarial_training_model_doc(
+        """Implementation on the training stage of AdversarialTrainingRegressor.""",  # noqa: E501
+        "fit"
+    )
+    def fit(self,
+            train_loader,
+            lr=1e-3,
+            weight_decay=5e-4,
+            epochs=100,
+            optimizer="Adam",
+            epsilon=0.5,
+            log_interval=100,
+            test_loader=None,
+            save_model=True,
+            save_dir=None):
+
+        # Instantiate base estimators and set attributes
+        estimators = []
+        for _ in range(self.n_estimators):
+            estimators.append(self._make_estimator())
+        self._validate_parameters(lr,
+                                  weight_decay,
+                                  epochs,
+                                  epsilon,
+                                  log_interval)
+        self.n_outputs = self._decide_n_outputs(train_loader, True)
+
+        # Utils
+        criterion = nn.MSELoss()
+        best_mse = float("inf")
+
+        # Internal helper function on pesudo forward
+        def _forward(estimators, data):
+            batch_size = data.size(0)
+            pred = torch.zeros(batch_size, self.n_outputs).to(self.device)
+
+            for estimator in estimators:
+                pred += estimator(data) / self.n_estimators
+
+            return pred
+
+        # Maintain a pool of workers
+        with Parallel(n_jobs=self.n_jobs) as parallel:
+
+            # Training loop
+            for epoch in range(epochs):
+                self.train()
+                rets = parallel(delayed(_parallel_fit_per_epoch)(
+                        train_loader,
+                        lr,
+                        weight_decay,
+                        epoch,
+                        optimizer,
+                        epsilon,
+                        log_interval,
+                        idx,
+                        estimator,
+                        criterion,
+                        self.device,
+                        False
+                    )
+                    for idx, estimator in enumerate(estimators)
+                )
+                estimators = rets  # update
+
+                # Validation
+                if test_loader:
+                    self.eval()
+                    with torch.no_grad():
+                        mse = 0
+                        for _, (data, target) in enumerate(test_loader):
+                            data = data.to(self.device)
+                            target = target.to(self.device)
+                            output = _forward(estimators, data)
+                            mse += criterion(output, target)
+                        mse /= len(test_loader)
+
+                        if mse < best_mse:
+                            best_mse = mse
+                            self.estimators_ = nn.ModuleList()
+                            self.estimators_.extend(estimators)
+                            if save_model:
+                                utils.save(self, save_dir, self.logger)
+
+                        msg = ("Epoch: {:03d} | Validation MSE:"
+                               " {:.5f} | Historical Best: {:.5f}")
+                        self.logger.info(msg.format(epoch, mse, best_mse))
+
+        self.estimators_ = nn.ModuleList()
+        self.estimators_.extend(rets)
+        if save_model and not test_loader:
+            utils.save(self, save_dir, self.logger)
+
+    @torchensemble_model_doc(
+        """Implementation on the evaluating stage of AdversarialTrainingRegressor.""",  # noqa: E501
+        "regressor_predict")
+    def predict(self, test_loader):
+        self.eval()
+        mse = 0
+        criterion = nn.MSELoss()
+
+        for batch_idx, (data, target) in enumerate(test_loader):
+            data, target = data.to(self.device), target.to(self.device)
+            output = self.forward(data)
+            mse += criterion(output, target)
+
+        return mse / len(test_loader)


Test patch:
diff --git a/torchensemble/tests/test_adversarial_training.py b/torchensemble/tests/test_adversarial_training.py
new file mode 100644
index 0000000..4859c36
--- /dev/null
+++ b/torchensemble/tests/test_adversarial_training.py
@@ -0,0 +1,50 @@
+import torch
+import pytest
+import numpy as np
+import torch.nn as nn
+from torch.utils.data import TensorDataset, DataLoader
+
+import torchensemble
+from torchensemble.utils import set_logger
+
+
+set_logger("pytest_adversarial_training")
+
+
+X_train = torch.Tensor(np.array(([1, 1],
+                                 [2, 2],
+                                 [3, 3],
+                                 [4, 4])))
+y_train_clf = torch.LongTensor(np.array(([0, 0, 1, 1])))
+
+
+# Base estimator
+class MLP_clf(nn.Module):
+    def __init__(self):
+        super(MLP_clf, self).__init__()
+        self.linear1 = nn.Linear(2, 2)
+        self.linear2 = nn.Linear(2, 2)
+
+    def forward(self, X):
+        X = X.view(X.size()[0], -1)
+        output = self.linear1(X)
+        output = self.linear2(output)
+        return output
+
+
+def test_adversarial_training_range():
+    """
+    This unit test checks the input range check in adversarial_training.py.
+    """
+    model = torchensemble.AdversarialTrainingClassifier(estimator=MLP_clf,
+                                                        n_estimators=2,
+                                                        cuda=False)
+
+    # Prepare data
+    train = TensorDataset(X_train, y_train_clf)
+    train_loader = DataLoader(train, batch_size=2)
+
+    # Training
+    with pytest.raises(ValueError) as excinfo:
+        model.fit(train_loader)
+    assert "input range of samples passed to adversarial" in str(excinfo.value)
diff --git a/torchensemble/tests/test_all_models.py b/torchensemble/tests/test_all_models.py
index 589d5c3..3eb6a36 100644
--- a/torchensemble/tests/test_all_models.py
+++ b/torchensemble/tests/test_all_models.py
@@ -12,14 +12,16 @@
            torchensemble.VotingClassifier,
            torchensemble.BaggingClassifier,
            torchensemble.GradientBoostingClassifier,
-           torchensemble.SnapshotEnsembleClassifier]
+           torchensemble.SnapshotEnsembleClassifier,
+           torchensemble.AdversarialTrainingClassifier]
 
 
 all_reg = [torchensemble.FusionRegressor,
            torchensemble.VotingRegressor,
            torchensemble.BaggingRegressor,
            torchensemble.GradientBoostingRegressor,
-           torchensemble.SnapshotEnsembleRegressor]
+           torchensemble.SnapshotEnsembleRegressor,
+           torchensemble.AdversarialTrainingRegressor]
 
 
 set_logger("pytest_all_models")
@@ -53,20 +55,20 @@ def forward(self, X):
 
 
 # Trainining data
-X_train = torch.Tensor(np.array(([1, 1],
-                                 [2, 2],
-                                 [3, 3],
-                                 [4, 4])))
+X_train = torch.Tensor(np.array(([0.1, 0.1],
+                                 [0.2, 0.2],
+                                 [0.3, 0.3],
+                                 [0.4, 0.4])))
 
 y_train_clf = torch.LongTensor(np.array(([0, 0, 1, 1])))
 y_train_reg = torch.FloatTensor(np.array(([0.1, 0.2, 0.3, 0.4])))
 y_train_reg = y_train_reg.view(-1, 1)
 
 # Testing data
-X_test = torch.Tensor(np.array(([5, 5],
-                                [6, 6],
-                                [7, 7],
-                                [8, 8])))
+X_test = torch.Tensor(np.array(([0.5, 0.5],
+                                [0.6, 0.6],
+                                [0.7, 0.7],
+                                [0.8, 0.8])))
 
 y_test_clf = torch.LongTensor(np.array(([1, 1, 0, 0])))
 y_test_reg = torch.FloatTensor(np.array(([0.5, 0.6, 0.7, 0.8])))
diff --git a/torchensemble/tests/test_training_params.py b/torchensemble/tests/test_training_params.py
index 5eac3d0..1787fd3 100644
--- a/torchensemble/tests/test_training_params.py
+++ b/torchensemble/tests/test_training_params.py
@@ -149,3 +149,34 @@ def test_snapshot_ensemble():
     with pytest.raises(ValueError) as excinfo:
         model.fit(train_loader, epochs=5)
     assert "should be a multiple of n_estimators" in str(excinfo.value)
+
+
+def test_adversarial_training():
+    model = torchensemble.AdversarialTrainingClassifier(estimator=MLP,
+                                                        n_estimators=2,
+                                                        cuda=False)
+
+    # Learning rate
+    with pytest.raises(ValueError) as excinfo:
+        model.fit(train_loader, lr=-1)
+    assert "learning rate of optimizer" in str(excinfo.value)
+
+    # Weight decay
+    with pytest.raises(ValueError) as excinfo:
+        model.fit(train_loader, weight_decay=-1)
+    assert "weight decay of optimizer" in str(excinfo.value)
+
+    # Epochs
+    with pytest.raises(ValueError) as excinfo:
+        model.fit(train_loader, epochs=-1)
+    assert "number of training epochs" in str(excinfo.value)
+
+    # Epsilon
+    with pytest.raises(ValueError) as excinfo:
+        model.fit(train_loader, epsilon=2)
+    assert "step used to generate adversarial samples" in str(excinfo.value)
+
+    # Log interval
+    with pytest.raises(ValueError) as excinfo:
+        model.fit(train_loader, log_interval=-1)
+    assert "number of batches to wait" in str(excinfo.value)