Skip to content

Commit

Permalink
mypy fixes
Browse files Browse the repository at this point in the history
Signed-off-by: GiulioZizzo <giulio.zizzo@yahoo.co.uk>
  • Loading branch information
GiulioZizzo committed Dec 20, 2023
1 parent 56f7f4a commit 63916b2
Show file tree
Hide file tree
Showing 3 changed files with 10 additions and 6 deletions.
7 changes: 5 additions & 2 deletions art/estimators/certification/randomized_smoothing/pytorch.py
Original file line number Diff line number Diff line change
Expand Up @@ -222,17 +222,20 @@ def fit( # pylint: disable=W0221
if scheduler is not None:
scheduler.step()

def predict(self, x: np.ndarray, batch_size: int = 128, **kwargs) -> np.ndarray: # type: ignore
def predict(self, x: np.ndarray, batch_size: int = 128, verbose: bool = False, **kwargs) -> np.ndarray: # type: ignore
"""
Perform prediction of the given classifier for a batch of inputs, taking an expectation over transformations.
:param x: Input samples.
:param batch_size: Batch size.
:param verbose: Display training progress bar.
:param is_abstain: True if function will abstain from prediction and return 0s. Default: True
:type is_abstain: `boolean`
:return: Array of predictions of shape `(nb_inputs, nb_classes)`.
"""
return RandomizedSmoothingMixin.predict(self, x, batch_size=batch_size, training_mode=False, **kwargs)
return RandomizedSmoothingMixin.predict(
self, x, batch_size=batch_size, verbose=verbose, training_mode=False, **kwargs
)

def loss_gradient( # type: ignore
self, x: np.ndarray, y: np.ndarray, training_mode: bool = False, **kwargs
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,6 @@ def __init__(
sample_size=sample_size,
scale=scale,
alpha=alpha,
verbose=verbose,
)
self.epsilon = epsilon
self.num_noise_vec = num_noise_vec
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,6 @@ def __init__(
sample_size: int = 32,
scale: float = 0.1,
alpha: float = 0.001,
verbose: bool = False,
):
"""
Create a randomized smoothing classifier.
Expand Down Expand Up @@ -192,17 +191,20 @@ def train_step(model, images, labels):
if scheduler is not None:
scheduler(epoch)

def predict(self, x: np.ndarray, batch_size: int = 128, **kwargs) -> np.ndarray: # type: ignore
def predict(self, x: np.ndarray, batch_size: int = 128, verbose: bool = False, **kwargs) -> np.ndarray: # type: ignore
"""
Perform prediction of the given classifier for a batch of inputs, taking an expectation over transformations.
:param x: Input samples.
:param batch_size: Batch size.
:param verbose: Display training progress bar.
:param is_abstain: True if function will abstain from prediction and return 0s. Default: True
:type is_abstain: `boolean`
:return: Array of predictions of shape `(nb_inputs, nb_classes)`.
"""
return RandomizedSmoothingMixin.predict(self, x, batch_size=batch_size, training_mode=False, **kwargs)
return RandomizedSmoothingMixin.predict(
self, x, batch_size=batch_size, verbose=verbose, training_mode=False, **kwargs
)

def loss_gradient(self, x: np.ndarray, y: np.ndarray, training_mode: bool = False, **kwargs) -> np.ndarray:
"""
Expand Down

0 comments on commit 63916b2

Please sign in to comment.