shok.utils package

Subpackages

• shok.utils.functions package
  • PassRound
    • PassRound.backward()
    • PassRound.forward()
  • ScaleGrad
    • ScaleGrad.backward()
    • ScaleGrad.forward()
  • SoftRound
    • SoftRound.backward()
    • SoftRound.forward()
  • Submodules
  • shok.utils.functions.pass_round module
    • PassRound
      • PassRound.backward()
      • PassRound.forward()
  • shok.utils.functions.scale_grad module
    • ScaleGrad
      • ScaleGrad.backward()
      • ScaleGrad.forward()
  • shok.utils.functions.soft_round module
    • SoftRound
      • SoftRound.backward()
      • SoftRound.forward()
• shok.utils.transforms package
  • ApplyPatch
    • ApplyPatch.forward()
  • ConvertToTVTensorBBoxes
    • ConvertToTVTensorBBoxes.forward()
  • PassRound
    • PassRound.forward()
  • ScaleApplyPatch
    • ScaleApplyPatch.forward()
  • ScaleGradTransform
    • ScaleGradTransform.forward()
  • ScaleImageValues
    • ScaleImageValues.forward()
  • SoftRound
    • SoftRound.forward()
  • TargetInsurance
    • TargetInsurance.forward()
  • Submodules
  • shok.utils.transforms.apply_patch module
    • ApplyPatch
      • ApplyPatch.forward()
  • shok.utils.transforms.convert_to_tv_tensor_bboxes module
    • ConvertToTVTensorBBoxes
      • ConvertToTVTensorBBoxes.forward()
  • shok.utils.transforms.pass_round module
    • PassRound
      • PassRound.forward()
  • shok.utils.transforms.scale_apply_patch module
    • ScaleApplyPatch
      • ScaleApplyPatch.forward()
  • shok.utils.transforms.scale_grad_transform module
    • ScaleGradTransform
      • ScaleGradTransform.forward()
  • shok.utils.transforms.scale_image_values module
    • ScaleImageValues
      • ScaleImageValues.forward()
  • shok.utils.transforms.simple_apply_patch module
    • SimpleApplyPatch
      • SimpleApplyPatch.forward()
  • shok.utils.transforms.soft_round module
    • SoftRound
      • SoftRound.forward()
  • shok.utils.transforms.target_insurance module
    • TargetInsurance
      • TargetInsurance.forward()
  • shok.utils.transforms.utils module
    • default_patched_image_mutator()

Submodules

shok.utils.callbacks module

class shok.utils.callbacks.LogPerformanceCallback

Bases: Callback

Callback to log performance metrics during model training using PyTorch Lightning.

This callback tracks and logs the following timing metrics: - Total fit time (from start to stop of the fit process) - Total training time (from start to stop of training) - Training epoch time (duration of each training epoch)

Metrics are logged to the LightningModule using the log method, with the following keys: - “performance/fit_time” - “performance/train_time” - “performance/train_epoch_time”

All logging is performed only on rank zero in distributed settings.

Methods:

on_fit_start(trainer, pl_module): Records the start time of the fit process. on_fit_stop(trainer, pl_module): Logs the total fit time. on_train_start(trainer, pl_module): Records the start time of training. on_train_stop(trainer, pl_module): Logs the total training time. on_train_epoch_start(trainer, pl_module): Records the start time of the training epoch. on_train_epoch_end(trainer, pl_module): Logs the duration of the training epoch.

on_fit_start(trainer: Trainer, pl_module: LightningModule)

Called at the beginning of the model fitting process.

Args:

trainer (Trainer): The Trainer instance managing the training loop. pl_module (L.LightningModule): The LightningModule being trained.

Side Effects:

Initializes the fit_time attribute with the current time to track the duration of the fitting process.

on_fit_stop(trainer: Trainer, pl_module: LightningModule)

Called when the fitting process stops.

Args:

trainer (Trainer): The PyTorch Lightning Trainer instance managing the training loop. pl_module (L.LightningModule): The LightningModule being trained.

Logs the total fitting time under the metric name “performance/fit_time”.

on_train_epoch_end(trainer, pl_module)

Called at the end of each training epoch.

Args:

trainer: The PyTorch Lightning Trainer instance managing the training loop. pl_module: The LightningModule being trained.

Calculates the duration of the training epoch and logs it under the key “performance/train_epoch_time” for monitoring purposes.

on_train_epoch_start(trainer, pl_module)

Called at the start of each training epoch.

Args:

trainer: The PyTorch Lightning Trainer instance managing the training process. pl_module: The LightningModule being trained.

Side Effects:

Records the current time as the start time of the training epoch in self.train_epoch_start_time.

on_train_start(trainer: Trainer, pl_module: LightningModule)

Called at the beginning of the training process.

Args:

trainer (Trainer): The PyTorch Lightning Trainer instance managing the training loop. pl_module (L.LightningModule): The LightningModule being trained.

Side Effects:

Initializes and stores the start time of training in self.start_time.

on_train_stop(trainer: Trainer, pl_module: LightningModule)

Called when the training process stops.

Args:

trainer (Trainer): The PyTorch Lightning Trainer instance managing the training loop. pl_module (L.LightningModule): The LightningModule being trained.

Logs the total training time under the key “performance/train_time” at the end of the epoch.

class shok.utils.callbacks.WandbObjectDetectionCallback(patch_log_frequency: int = 1, train_log_frequency: int = 1, val_log_frequency: int = 1)

Bases: Callback

WandbObjectDetectionCallback.

A PyTorch Lightning callback for logging object detection images and bounding boxes to Weights & Biases (WandB) during training and validation.

This callback supports logging images with bounding boxes at configurable frequencies for both training and validation phases. It automatically extracts class information from the provided datamodule and formats it for WandB visualization.

patch_log_frequency (int): Frequency (in epochs) to log patch images at the end of training epochs. train_log_frequency (int): Frequency (in epochs) to log training images and bounding boxes. val_log_frequency (int): Frequency (in epochs) to log validation images and bounding boxes.

Raises:

ValueError: If WandbLogger is not set up or if the datamodule is not configured.

Methods:

setup(trainer, pl_module, stage=None): Initializes WandB experiment and class mapping from the datamodule. on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0): Logs training images and bounding boxes at the specified frequency. on_validation_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0): Logs validation images and bounding boxes at the specified frequency. on_train_epoch_end(trainer, pl_module): Logs the patch image at the end of the training epoch if available.

Note:

• The callback expects the datamodule to have an idx_to_class attribute for mapping class indices to names.

• Images and bounding boxes are logged using WandB’s image and box utilities for visualization.

on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0)

Called at the end of each training batch during model training.

Logs input images and their corresponding labels to Weights & Biases (WandB) at specified intervals, determined by train_log_frequency. Only logs when batch_idx is 0 and the current epoch is a multiple of train_log_frequency.

Args:

trainer: The PyTorch Lightning trainer instance. pl_module: The LightningModule being trained. outputs: The outputs from the training step. batch: The current batch of data, typically a tuple (x, y) of inputs and labels. batch_idx (int): Index of the current batch within the epoch. dataloader_idx (int, optional): Index of the dataloader, defaults to 0.

Notes:

• The method can be extended to mutate input images using model attributes if available.

• Images and labels are logged using the _wandb_log_image helper method.

on_train_epoch_end(trainer, pl_module)

Called at the end of each training epoch.

Logs the patch image to Weights & Biases (wandb) at intervals defined by self.train_log_frequency. If the pl_module has a patch attribute, the patch image is logged using wandb.Image. If the patch attribute is not present, no image is logged.

Args:

trainer: The PyTorch Lightning trainer instance. pl_module: The PyTorch Lightning module being trained.

Note:

Logging occurs only if the current epoch is divisible by self.train_log_frequency.

on_validation_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0)

Called at the end of each validation batch during model training.

Logs validation patch images to Weights & Biases (WandB) at the start of each epoch, based on the specified validation log frequency.

Args:

trainer: The PyTorch Lightning Trainer instance. pl_module: The LightningModule being trained. outputs: The outputs from the validation step. batch: The current batch of data, typically a tuple (x, y). batch_idx (int): Index of the current batch within the epoch. dataloader_idx (int, optional): Index of the dataloader. Defaults to 0.

setup(trainer, pl_module, stage=None)

Sets up the callback by initializing the Wandb logger and extracting class information from the datamodule.

Args:

trainer: The PyTorch Lightning trainer instance. pl_module: The PyTorch Lightning module (model) being trained. stage (optional): The current stage of training (e.g., ‘fit’, ‘test’).

Raises:

ValueError: If the Wandb logger is not set up or if the datamodule is not configured.

Side Effects:

• Initializes self._wandb with the Wandb experiment object.

• Sets self.idx_to_class from the datamodule if available.

• Creates self.wandb_classes as a wandb.Classes object if class mapping is provided.

shok.utils.callbacks.translate_pred_to_wandb_boxes(preds, patch_image)

Converts prediction results into a format compatible with Weights & Biases (wandb) for bounding box visualization.

Args:

preds (dict): A dictionary containing prediction results with keys “boxes”, “scores”, and “labels”.

• “boxes” (list or tensor): Bounding box coordinates.

• “scores” (list or tensor, optional): Confidence scores for each box. If missing, defaults to 100.0.

• “labels” (list or tensor): Class labels for each box.

patch_image: The image associated with the predictions, used for box translation.

Returns:

dict: A dictionary with a “predictions” key containing “box_data”, which is a list of bounding boxes formatted for wandb visualization.

shok.utils.callbacks.translate_pytorch_boxes_to_wandb(boxes, score, label, patch_image)

Converts PyTorch bounding box coordinates to the format expected by Weights & Biases (wandb).

Args:

boxes (torch.Tensor): A tensor containing bounding box coordinates in the format [minX, minY, maxX, maxY]. score (torch.Tensor or None): A tensor containing the confidence score for the bounding box, or None. label (torch.Tensor): A tensor containing the class label for the bounding box. patch_image (torch.Tensor): The image tensor corresponding to the bounding box, used to normalize coordinates.

Returns:

dict: A dictionary containing the normalized bounding box position, class ID, and score, formatted for wandb.

shok.utils.optim module

class shok.utils.optim.SubPixelOptimizer(params, lr=20.0)

Bases: Optimizer

A custom optimizer that operates on the pixel level.

step(closure=None)

Performs a single optimization step.

If a closure is provided, it is called and its result is returned as the loss.

For each parameter group and each parameter within the group:

• If the parameter has a gradient, normalize the gradient by its maximum absolute value (if non-zero).

• Update the parameter data by adding the scaled gradient multiplied by the learning rate.

Args:

closure (callable, optional): A closure that reevaluates the model and returns the loss.

Returns:

The loss value returned by the closure, if provided; otherwise, None.

class shok.utils.optim.WholePixelOptimizer(params, lr=0.01)

Bases: Optimizer

A custom optimizer that operates on the pixel level by stepping by whole pixel values.

This approach is common in adversarial patch training.

step(closure=None)

Performs a single optimization step.

This method updates the parameters by adding the sign of their gradients multiplied by a rounded learning rate. The learning rate is rounded to the nearest integer and is at least 1.0.

Args:

closure (callable, optional): A closure that reevaluates the model and returns the loss.

Returns:

loss: The loss value returned by the closure, if provided; otherwise, None.