robustness.model_utils module

class robustness.model_utils.FeatureExtractor(submod, layers)

Bases: sphinx.ext.autodoc.importer._MockObject

Tool for extracting layers from models.

  • submod (torch.nn.Module) – model to extract activations from
  • layers (list of functions) – list of functions where each function, when applied to submod, returns a desired layer. For example, one function could be lambda model: model.layer1.

A model whose forward function returns the activations from the layers

corresponding to the functions in layers (in the order that the functions were passed in the list).

forward(*args, **kwargs)
robustness.model_utils.make_and_restore_model(*_, arch, dataset, resume_path=None, parallel=True, pytorch_pretrained=False)

Makes a model and (optionally) restores it from a checkpoint.

  • arch (str|nn.Module) – Model architecture identifier or otherwise a torch.nn.Module instance with the classifier
  • dataset (Dataset class [see]) –
  • resume_path (str) – optional path to checkpoint
  • parallel (bool) – if True, wrap the model in a DataParallel (default True, recommended)
  • pytorch_pretrained (bool) – if True, try to load a standard-trained checkpoint from the torchvision library (throw error if failed)

A tuple consisting of the model (possibly loaded with checkpoint), and the checkpoint itself

robustness.model_utils.model_dataset_from_store(s, overwrite_params={}, which='last')

Given a store directory corresponding to a trained model, return the original model, dataset object, and args corresponding to the arguments.