├── .gitignore ├── .gitmodules ├── LICENSE.txt ├── MANIFEST.in ├── README.rst ├── docs ├── Makefile ├── _static │ └── all.min.css ├── api.rst ├── api │ ├── robustness.attack_steps.rst │ ├── robustness.attacker.rst │ ├── robustness.data_augmentation.rst │ ├── robustness.datasets.rst │ ├── robustness.defaults.rst │ ├── robustness.loaders.rst │ ├── robustness.main.rst │ ├── robustness.model_utils.rst │ ├── robustness.tools.breeds_helpers.rst │ ├── robustness.tools.constants.rst │ ├── robustness.tools.folder.rst │ ├── robustness.tools.helpers.rst │ ├── robustness.tools.label_maps.rst │ ├── robustness.tools.rst │ ├── robustness.tools.vis_tools.rst │ └── robustness.train.rst ├── conf.py ├── example_usage │ ├── Figures │ │ ├── breeds_pipeline.png │ │ ├── breeds_superclasses.png │ │ ├── custom_inversion_CIFAR.png │ │ ├── custom_inversion_CIFAR_fourier.png │ │ ├── targeted_adversarial_example_CIFAR.png │ │ └── untargeted_adversarial_example_CIFAR.png │ ├── breeds_datasets.rst │ ├── changelog.rst │ ├── cli_usage.rst │ ├── custom_imagenet.rst │ ├── input_space_manipulation.rst │ ├── training_lib_part_1.rst │ └── training_lib_part_2.rst ├── index.rst ├── make.bat └── requirements.txt ├── notebooks ├── Input manipulation with pre-trained models.ipynb ├── Superclassing ImageNet.ipynb └── Using robustness as a library.ipynb ├── requirements.txt ├── robustness ├── .gitignore ├── __init__.py ├── attack_steps.py ├── attacker.py ├── cifar_models │ ├── __init__.py │ ├── densenet.py │ ├── inception.py │ ├── resnet.py │ └── vgg.py ├── data_augmentation.py ├── datasets.py ├── defaults.py ├── imagenet_models │ ├── __init__.py │ ├── alexnet.py │ ├── densenet.py │ ├── leaky_resnet.py │ ├── resnet.py │ ├── squeezenet.py │ └── vgg.py ├── loaders.py ├── main.py ├── model_utils.py ├── tools │ ├── __init__.py │ ├── breeds_helpers.py │ ├── constants.py │ ├── custom_modules.py │ ├── folder.py │ ├── helpers.py │ ├── imagenet_helpers.py │ ├── label_maps.py │ ├── openimgs_helpers.py │ └── vis_tools.py └── train.py └── setup.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | .vscode 6 | 7 | # C extensions 8 | *.so 9 | 10 | # Distribution / packaging 11 | .Python 12 | build/ 13 | develop-eggs/ 14 | dist/ 15 | downloads/ 16 | eggs/ 17 | .eggs/ 18 | lib/ 19 | lib64/ 20 | parts/ 21 | sdist/ 22 | var/ 23 | wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | MANIFEST 28 | 29 | # PyInstaller 30 | # Usually these files are written by a python script from a template 31 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 32 | *.manifest 33 | *.spec 34 | 35 | # Installer logs 36 | pip-log.txt 37 | pip-delete-this-directory.txt 38 | 39 | # Unit test / coverage reports 40 | htmlcov/ 41 | .tox/ 42 | .coverage 43 | .coverage.* 44 | .cache 45 | nosetests.xml 46 | coverage.xml 47 | *.cover 48 | .hypothesis/ 49 | .pytest_cache/ 50 | 51 | # Translations 52 | *.mo 53 | *.pot 54 | 55 | # Django stuff: 56 | *.log 57 | local_settings.py 58 | db.sqlite3 59 | 60 | # Flask stuff: 61 | instance/ 62 | .webassets-cache 63 | 64 | # Scrapy stuff: 65 | .scrapy 66 | 67 | # Sphinx documentation 68 | docs/_build/ 69 | 70 | # PyBuilder 71 | target/ 72 | 73 | # Jupyter Notebook 74 | .ipynb_checkpoints 75 | 76 | # pyenv 77 | .python-version 78 | 79 | # celery beat schedule file 80 | celerybeat-schedule 81 | 82 | # SageMath parsed files 83 | *.sage.py 84 | 85 | # Environments 86 | .env 87 | .venv 88 | env/ 89 | venv/ 90 | ENV/ 91 | env.bak/ 92 | venv.bak/ 93 | 94 | # Spyder project settings 95 | .spyderproject 96 | .spyproject 97 | 98 | # Rope project settings 99 | .ropeproject 100 | 101 | # mkdocs documentation 102 | /site 103 | 104 | # mypy 105 | .mypy_cache/ 106 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "src/cox"] 2 | path = src/cox 3 | url = git@github.com:andrewilyas/cox.git 4 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 Andrew Ilyas, Logan Engstrom 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include LICENSE.txt 2 | -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- 1 | # Minimal makefile for Sphinx documentation 2 | # 3 | 4 | # You can set these variables from the command line, and also 5 | # from the environment for the first two. 6 | SPHINXOPTS ?= 7 | SPHINXBUILD ?= sphinx-build 8 | SOURCEDIR = . 9 | BUILDDIR = _build 10 | 11 | # Put it first so that "make" without argument is like "make help". 12 | help: 13 | @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 14 | 15 | .PHONY: help Makefile 16 | 17 | # Catch-all target: route all unknown targets to Sphinx using the new 18 | # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). 19 | %: Makefile 20 | @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 21 | -------------------------------------------------------------------------------- /docs/api.rst: -------------------------------------------------------------------------------- 1 | API Reference 2 | ============= 3 | 4 | .. toctree:: 5 | api/robustness.attack_steps 6 | api/robustness.attacker 7 | api/robustness.data_augmentation 8 | api/robustness.datasets 9 | api/robustness.defaults 10 | api/robustness.loaders 11 | api/robustness.main 12 | api/robustness.model_utils 13 | api/robustness.train 14 | api/robustness.tools 15 | 16 | -------------------------------------------------------------------------------- /docs/api/robustness.attack_steps.rst: -------------------------------------------------------------------------------- 1 | robustness.attack\_steps module 2 | =============================== 3 | 4 | .. automodule:: robustness.attack_steps 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /docs/api/robustness.attacker.rst: -------------------------------------------------------------------------------- 1 | robustness.attacker module 2 | ========================== 3 | 4 | .. automodule:: robustness.attacker 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /docs/api/robustness.data_augmentation.rst: -------------------------------------------------------------------------------- 1 | robustness.data\_augmentation module 2 | ==================================== 3 | 4 | .. automodule:: robustness.data_augmentation 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /docs/api/robustness.datasets.rst: -------------------------------------------------------------------------------- 1 | robustness.datasets module 2 | ========================== 3 | 4 | .. automodule:: robustness.datasets 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /docs/api/robustness.defaults.rst: -------------------------------------------------------------------------------- 1 | robustness.defaults module 2 | ========================== 3 | 4 | .. automodule:: robustness.defaults 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /docs/api/robustness.loaders.rst: -------------------------------------------------------------------------------- 1 | robustness.loaders module 2 | ========================= 3 | 4 | .. automodule:: robustness.loaders 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /docs/api/robustness.main.rst: -------------------------------------------------------------------------------- 1 | robustness.main module 2 | ====================== 3 | 4 | .. automodule:: robustness.main 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /docs/api/robustness.model_utils.rst: -------------------------------------------------------------------------------- 1 | robustness.model\_utils module 2 | ============================== 3 | 4 | .. automodule:: robustness.model_utils 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /docs/api/robustness.tools.breeds_helpers.rst: -------------------------------------------------------------------------------- 1 | robustness.tools.breeds\_helpers module 2 | ================================== 3 | 4 | .. automodule:: robustness.tools.breeds_helpers 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /docs/api/robustness.tools.constants.rst: -------------------------------------------------------------------------------- 1 | robustness.tools.constants module 2 | ================================= 3 | 4 | .. automodule:: robustness.tools.constants 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /docs/api/robustness.tools.folder.rst: -------------------------------------------------------------------------------- 1 | robustness.tools.folder module 2 | ============================== 3 | 4 | .. automodule:: robustness.tools.folder 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /docs/api/robustness.tools.helpers.rst: -------------------------------------------------------------------------------- 1 | robustness.tools.helpers module 2 | =============================== 3 | 4 | .. automodule:: robustness.tools.helpers 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /docs/api/robustness.tools.label_maps.rst: -------------------------------------------------------------------------------- 1 | robustness.tools.label\_maps module 2 | =================================== 3 | 4 | .. automodule:: robustness.tools.label_maps 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /docs/api/robustness.tools.rst: -------------------------------------------------------------------------------- 1 | robustness.tools package 2 | ======================== 3 | 4 | Submodules 5 | ---------- 6 | 7 | .. toctree:: 8 | 9 | robustness.tools.constants 10 | robustness.tools.folder 11 | robustness.tools.helpers 12 | robustness.tools.label_maps 13 | robustness.tools.vis_tools 14 | robustness.tools.breeds_helpers 15 | 16 | Module contents 17 | --------------- 18 | 19 | .. automodule:: robustness.tools 20 | :members: 21 | :undoc-members: 22 | :show-inheritance: 23 | -------------------------------------------------------------------------------- /docs/api/robustness.tools.vis_tools.rst: -------------------------------------------------------------------------------- 1 | robustness.tools.vis\_tools module 2 | ================================== 3 | 4 | .. automodule:: robustness.tools.vis_tools 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /docs/api/robustness.train.rst: -------------------------------------------------------------------------------- 1 | robustness.train module 2 | ======================= 3 | 4 | .. automodule:: robustness.train 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /docs/conf.py: -------------------------------------------------------------------------------- 1 | # Configuration file for the Sphinx documentation builder. 2 | # 3 | # This file only contains a selection of the most common options. For a full 4 | # list see the documentation: 5 | # http://www.sphinx-doc.org/en/master/config 6 | 7 | # -- Path setup -------------------------------------------------------------- 8 | 9 | # If extensions (or modules to document with autodoc) are in another directory, 10 | # add these directories to sys.path here. If the directory is relative to the 11 | # documentation root, use os.path.abspath to make it absolute, like shown here. 12 | # 13 | import os 14 | import sys 15 | sys.path.insert(0, os.path.abspath('..')) 16 | 17 | 18 | # -- Project information ----------------------------------------------------- 19 | 20 | project = 'robustness' 21 | copyright = '2019, MadryLab' 22 | author = 'MadryLab' 23 | 24 | # The full version, including alpha/beta/rc tags 25 | release = '1.0' 26 | 27 | 28 | # -- General configuration --------------------------------------------------- 29 | 30 | # Add any Sphinx extension module names here, as strings. They can be 31 | # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom 32 | # ones. 33 | extensions = [ 34 | 'sphinx.ext.autodoc', 35 | 'sphinx.ext.autosummary', 36 | 'sphinx.ext.napoleon', 37 | 'sphinx.ext.mathjax' 38 | ] 39 | 40 | napoleon_google_docstring = True 41 | 42 | autodoc_default_options = { 43 | 'undoc-members': False 44 | } 45 | 46 | autodoc_mock_imports = ['torch', 'numpy', 'pandas', 'tensorboardX', 'dill', 47 | 'torchvision', 'scipy', 'GPUtil', 'tables', 'scikit-learn', 'seaborn', 48 | 'cox', 'matplotlib', 'sklearn'] 49 | 50 | autoclass_content = "both" 51 | 52 | # Add any paths that contain templates here, relative to this directory. 53 | templates_path = ['_templates'] 54 | 55 | # List of patterns, relative to source directory, that match files and 56 | # directories to ignore when looking for source files. 57 | # This pattern also affects html_static_path and html_extra_path. 58 | exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store', 59 | '**/imagenet_models', '**/cifar_models'] 60 | 61 | 62 | # -- Options for HTML output ------------------------------------------------- 63 | 64 | # The theme to use for HTML and HTML Help pages. See the documentation for 65 | # a list of builtin themes. 66 | # 67 | html_theme = 'sphinx_rtd_theme' 68 | html_theme_options = { 69 | 'collapse_navigation': False, 70 | 'navigation_depth': -1, 71 | 'includehidden': True, 72 | } 73 | 74 | master_doc = 'index' 75 | 76 | # Add any paths that contain custom static files (such as style sheets) here, 77 | # relative to this directory. They are copied after the builtin static files, 78 | # so a file named "default.css" will overwrite the builtin "default.css". 79 | html_static_path = ['_static'] 80 | html_css_files = ['all.min.css'] 81 | 82 | autodoc_member_order = 'bysource' 83 | -------------------------------------------------------------------------------- /docs/example_usage/Figures/breeds_pipeline.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MadryLab/robustness/a9541241defd9972e9334bfcdb804f6aefe24dc7/docs/example_usage/Figures/breeds_pipeline.png -------------------------------------------------------------------------------- /docs/example_usage/Figures/breeds_superclasses.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MadryLab/robustness/a9541241defd9972e9334bfcdb804f6aefe24dc7/docs/example_usage/Figures/breeds_superclasses.png -------------------------------------------------------------------------------- /docs/example_usage/Figures/custom_inversion_CIFAR.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MadryLab/robustness/a9541241defd9972e9334bfcdb804f6aefe24dc7/docs/example_usage/Figures/custom_inversion_CIFAR.png -------------------------------------------------------------------------------- /docs/example_usage/Figures/custom_inversion_CIFAR_fourier.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MadryLab/robustness/a9541241defd9972e9334bfcdb804f6aefe24dc7/docs/example_usage/Figures/custom_inversion_CIFAR_fourier.png -------------------------------------------------------------------------------- /docs/example_usage/Figures/targeted_adversarial_example_CIFAR.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MadryLab/robustness/a9541241defd9972e9334bfcdb804f6aefe24dc7/docs/example_usage/Figures/targeted_adversarial_example_CIFAR.png -------------------------------------------------------------------------------- /docs/example_usage/Figures/untargeted_adversarial_example_CIFAR.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MadryLab/robustness/a9541241defd9972e9334bfcdb804f6aefe24dc7/docs/example_usage/Figures/untargeted_adversarial_example_CIFAR.png -------------------------------------------------------------------------------- /docs/example_usage/breeds_datasets.rst: -------------------------------------------------------------------------------- 1 | Creating BREEDS subpopulation shift benchmarks 2 | =============================================== 3 | 4 | In this document, we will discuss how to create BREEDS datasets [STM20]_. 5 | Given any existing dataset that comes with a class hierarchy (e.g. ImageNet, 6 | OpenImages), the BREEDS methodology allows you to make a derivative 7 | classification task that can be used to measure robustness to subpopulation 8 | shift. To do this, we: 9 | 10 | 1. Group together semantically-simlar classes ("breeds") in the dataset 11 | into superclasses. 12 | 2. Define a classification task in terms of these superclasses---with 13 | the twist that the "breeds" used in the training set from each superclasses 14 | are disjoint from the "breeds" used in the test set. 15 | 16 | As a primitive example, one could take ImageNet (which contains many classes 17 | corresponding to cat and dog breeds), and use the BREEDS methodology to come up 18 | with a derivative "cats vs. dogs" task, where the training set would contain one 19 | set of breeds (e.g., Egyptian cat and Tabby Cat vs. Labrador and Golden 20 | Retriever) and the test set would contain another set (e.g. Persian cat and 21 | alley cat vs Mastiff and Poodle). Here is a pictorial illustration of the BREEDS 22 | approach: 23 | 24 | .. image:: Figures/breeds_pipeline.png 25 | :width: 600 26 | :align: center 27 | :alt: Illustration of the BREEDS dataset creation pipeline. 28 | 29 | This methodology allows you to create subpopulation shift benchmarks of varying 30 | difficulty automatically, without having to manually group or split up classes, 31 | and can be applied to any dataset which has a class hierarchy. In this 32 | walkthrough, we will use ImageNet and the corresponding class hierarchy from 33 | [STM20]_. 34 | 35 | .. raw:: html 36 | 37 |    Download 39 | a Jupyter notebook containing all the code from this walkthrough!
40 |
41 | 42 | Requirements/Setup 43 | '''''''''''''''''' 44 | To create BREEDS datasets using ImageNet, we need to create a: 45 | 46 | - ``data_dir`` which contains the ImageNet dataset 47 | in PyTorch-readable format. 48 | - ``info_dir`` which contains the following information (files) about 49 | the class hierarchy: 50 | 51 | - ``dataset_class_info.json``: A list whose entries are triplets of 52 | class number, class ID and class name, for each dataset class. 53 | - ``class_hierarchy.txt``: Every line denotes an edge---parent ID followed by 54 | child ID (space separated)---in the class hierarchy. 55 | - ``node_names.txt``: Each line contains the ID of a node followed by 56 | it's name (tab separated). 57 | 58 | For convenience, we provide the relevant files for the (modified) class 59 | hierarchy `here 60 | `_. 61 | You can manually download them and move them to ``info_dir`` or do it 62 | automatically by specifying an empty ``info_dir`` to 63 | :meth:`~robustness.tools.breeds_helpers.BreedsDatasetGenerator.get_superclasses`: 64 | 65 | .. code-block:: python 66 | 67 | from robustness.tools.breeds_helpers import setup_breeds 68 | 69 | setup_breeds(info_dir) 70 | 71 | 72 | Part 1: Browsing through the Class Hierarchy 73 | '''''''''''''''''''''''''''''''''''''''''''' 74 | 75 | We can use :class:`~robustness.tools.breeds_helpers.ClassHierarchy` to 76 | examine a dataset's (here, ImageNet) class hierarchy. Here, ``info_dir`` 77 | should contain the requisite files for the class hierarchy (from the Setup 78 | step): 79 | 80 | .. code-block:: python 81 | 82 | from robustness.tools.breeds_helpers import ClassHierarchy 83 | import numpy as np 84 | 85 | hier = ClassHierarchy(info_dir) 86 | print(f"# Levels in hierarchy: {np.max(list(hier.level_to_nodes.keys()))}") 87 | print(f"# Nodes/level:", 88 | [f"Level {k}: {len(v)}" for k, v in hier.level_to_nodes.items()]) 89 | 90 | The :samp:`hier` object has a ``graph`` attribute, which represents the class 91 | hierarchy as a ``networkx`` graph. In this graph, the children of a node 92 | correspond to its subclasses (e.g., Labrador would be a child of the dog 93 | class in our primitive example). Note that all the original dataset classes 94 | will be the leaves of this graph. 95 | 96 | We can then use this graph to define superclasses---all nodes at a user-specified 97 | depth from the root node. For example: 98 | 99 | .. code-block:: python 100 | 101 | level = 2 # Could be any number smaller than max level 102 | superclasses = hier.get_nodes_at_level(level) 103 | print(f"Superclasses at level {level}:\n") 104 | print(", ".join([f"{hier.HIER_NODE_NAME[s]}" for s in superclasses])) 105 | 106 | Each superclass is made up of multiple "breeds", which simply correspond to 107 | the leaves (original dataset classes) that are its descendants in the class 108 | hierarchy: 109 | 110 | .. code-block:: python 111 | 112 | idx = np.random.randint(0, len(superclasses), 1)[0] 113 | superclass = list(superclasses)[idx] 114 | subclasses = hier.leaves_reachable(superclass) 115 | print(f"Superclass: {hier.HIER_NODE_NAME[superclass]}\n") 116 | 117 | print(f"Subclasses ({len(subclasses)}):") 118 | print([f"{hier.LEAF_ID_TO_NAME[l]}" for l in list(subclasses)]) 119 | 120 | 121 | We can also visualize subtrees of the graph with the help of 122 | the `networkx` and `pygraphviz` packages. For instance, we can 123 | taks a look at the subtree of the class hierarchy rooted at a 124 | particular superclass: 125 | 126 | .. code-block:: python 127 | 128 | import networkx as nx 129 | from networkx.drawing.nx_agraph import graphviz_layout, to_agraph 130 | import pygraphviz as pgv 131 | from IPython.display import Image 132 | 133 | subtree = nx.ego_graph(hier.graph, superclass, radius=10) 134 | mapping = {n: hier.HIER_NODE_NAME[n] for n in subtree.nodes()} 135 | subtree = to_agraph(nx.relabel_nodes(subtree, mapping)) 136 | subtree.delete_edge(subtree.edges()[0]) 137 | subtree.layout('dot') 138 | subtree.node_attr['color']='blue' 139 | subtree.draw('graph.png', format='png') 140 | Image('graph.png') 141 | 142 | For instance, visualizing tree rooted at the ``fungus`` superclass yields: 143 | 144 | .. image:: Figures/breeds_superclasses.png 145 | :width: 600 146 | :align: center 147 | :alt: Visulization of subtree rooted at a specific superclass. 148 | 149 | Part 2: Creating BREEDS Datasets 150 | ''''''''''''''''''''''''''''''''' 151 | 152 | To create a dataset composed of superclasses, we use the 153 | :class:`~robustness.tools.breeds_helpers.BreedsDatasetGenerator`. 154 | Internally, this class instantiates an object of 155 | :class:`~robustness.tools.breeds_helpers.ClassHierarchy` and uses it 156 | to define the superclasses. 157 | 158 | .. code-block:: python 159 | 160 | from robustness.tools.breeds_helpers import BreedsDatasetGenerator 161 | 162 | DG = BreedsDatasetGenerator(info_dir) 163 | 164 | Specifically, we will use 165 | :meth:`~robustness.tools.breeds_helpers.BreedsDatasetGenerator.get_superclasses`. 166 | This function takes in the following arguments (see :meth:`this docstring 167 | ` for more details): 168 | 169 | - :samp:`level`: Level in the hierarchy (in terms of distance from the 170 | root node) at which to define superclasses. 171 | - :samp:`Nsubclasses`: Controls the minimum number of subclasses/superclass 172 | in the dataset. If None, it is automatically set to be the size (in terms 173 | of subclasses) of the smallest superclass. 174 | - :samp:`split`: If ``None``, subclasses of a superclass are returned 175 | as is, without partitioning them into the source and target domains. 176 | Else, can be ``rand/good/bad`` depending on whether the subclass split should be 177 | random or less/more adversarially chosen (see paper for details). 178 | - :samp:`ancestor`: If a node ID is specified, superclasses are chosen from 179 | subtree of class hierarchy rooted at this node. Else, if None, :samp:`ancestor` 180 | is set to be the root node. 181 | - :samp:`balanced`: If True, subclasses/superclass is fixed over superclasses. 182 | 183 | For instance, we could create a balanced dataset, with the subclass partition 184 | being less adversarial as follows: 185 | 186 | .. code-block:: python 187 | 188 | ret = DG.get_superclasses(level=2, 189 | Nsubclasses=None, 190 | split="rand", 191 | ancestor=None, 192 | balanced=True) 193 | superclasses, subclass_split, label_map = ret 194 | 195 | This method returns: 196 | 197 | - :samp:`superclasses` is a list containing the IDs of all the 198 | superclasses. 199 | - :samp:`subclass_split` is a tuple of subclass ranges for 200 | the source and target domains. For instance, 201 | :samp:`subclass_split[0]` is a list, which for each superclass, 202 | contains a list of subclasses present in the source domain. 203 | If ``split=None``, subclass_split[1] is empty and can be 204 | ignored. 205 | - :samp:`label_map` is a dictionary mapping a superclass 206 | number (label) to name. 207 | 208 | You can experiment with these parameters to create datasets of different 209 | granularity. For instance, you could specify the :samp:`Nsubclasses` to 210 | restrict the size of every superclass in the dataset, 211 | set the :samp:`ancestor` to be a specific node (e.g., ``n00004258`` 212 | to focus on living things), or set :samp:`balanced` to ``False`` 213 | to get an imbalanced dataset. 214 | 215 | We can take a closer look at the composition of the dataset---what 216 | superclasses/subclasses it contains---using: 217 | 218 | .. code-block:: python 219 | 220 | from robustness.tools.breeds_helpers import print_dataset_info 221 | 222 | print_dataset_info(superclasses, 223 | subclass_split, 224 | label_map, 225 | hier.LEAF_NUM_TO_NAME) 226 | 227 | Finally, for the source and target domains, we can create datasets 228 | and their corresponding loaders: 229 | 230 | .. code-block:: python 231 | 232 | from robustness import datasets 233 | 234 | train_subclasses, test_subclasses = subclass_split 235 | 236 | dataset_source = datasets.CustomImageNet(data_dir, train_subclasses) 237 | loaders_source = dataset_source.make_loaders(num_workers, batch_size) 238 | train_loader_source, val_loader_source = loaders_source 239 | 240 | dataset_target = datasets.CustomImageNet(data_dir, test_subclasses) 241 | loaders_target = dataset_source.make_loaders(num_workers, batch_size) 242 | train_loader_target, val_loader_target = loaders_target 243 | 244 | You're all set! You can then use this dataset and loaders 245 | just as you would any other existing/custom dataset in the robustness 246 | library. For instance, you can visualize validation set samples from 247 | both domains and their labels using: 248 | 249 | .. code-block:: python 250 | 251 | from robustness.tools.vis_tools import show_image_row 252 | 253 | for domain, loader in zip(["Source", "Target"], 254 | [val_loader_source, val_loader_target]): 255 | im, lab = next(iter(loader)) 256 | show_image_row([im], 257 | tlist=[[label_map[int(k)].split(",")[0] for k in lab]], 258 | ylist=[domain], 259 | fontsize=20) 260 | 261 | You can also create superclass tasks where subclasses are not 262 | partitioned across domains: 263 | 264 | .. code-block:: python 265 | 266 | ret = DG.get_superclasses(level=2, 267 | Nsubclasses=2, 268 | split=None, 269 | ancestor=None, 270 | balanced=True) 271 | superclasses, subclass_split, label_map = ret 272 | all_subclasses = subclass_split[0] 273 | 274 | dataset = datasets.CustomImageNet(data_dir, all_subclasses) 275 | 276 | print_dataset_info(superclasses, 277 | subclass_split, 278 | label_map, 279 | hier.LEAF_NUM_TO_NAME) 280 | 281 | Part 3: Loading in-built BREEDS Datasets 282 | '''''''''''''''''''''''''''''''''''''''' 283 | 284 | Alternatively, we can directly use one of the datasets from our paper 285 | [STM20]_---namely ``Entity13``, ``Entity30``, ``Living17`` 286 | and ``Nonliving26``. Loading any of these datasets is relatively simple: 287 | 288 | .. code-block:: python 289 | 290 | from robustness.tools.breeds_helpers import make_living17 291 | ret = make_living17(info_dir, split="rand") 292 | superclasses, subclass_split, label_map = ret 293 | 294 | print_dataset_info(superclasses, 295 | subclass_split, 296 | label_map, 297 | hier.LEAF_NUM_TO_NAME) 298 | 299 | You can then use a similar methodology to Part 2 above to probe 300 | dataset information and create datasets and loaders. 301 | 302 | -------------------------------------------------------------------------------- /docs/example_usage/changelog.rst: -------------------------------------------------------------------------------- 1 | CHANGELOG 2 | ========= 3 | 4 | robustness 1.2.post2 5 | ''''''''''''''''''''' 6 | - Add SqueezeNet architectures 7 | - (Preliminary) Torch 1.7 support 8 | - Support for specifying device_ids in DataParallel 9 | 10 | robustness 1.2 11 | ''''''''''''''' 12 | - Biggest new features: 13 | - New ImageNet models 14 | - Mixed-precision training 15 | - OpenImages and Places365 datasets added 16 | - Ability to specify a custom accuracy function (custom loss functions 17 | were already supported, this is just for logging) 18 | - Improved resuming functionality 19 | - Changes to CLI-based training: 20 | - ``--custom-lr-schedule`` replaced by ``--custom-lr-multiplier`` (same format) 21 | - ``--eps-fadein-epochs`` replaced by general ``--custom-eps-multiplier`` 22 | (now same format as custom-lr schedule) 23 | - ``--step-lr-gamma`` now available to change the size of learning rate 24 | drops (used to be fixed to 10x drops) 25 | - ``--lr-interpolation`` argument added (can choose between linear and step 26 | interpolation between learning rates in the schedule) 27 | - ``--weight_decay`` is now called ``--weight-decay``, keeping with 28 | convention 29 | - ``--resume-optimizer`` is a 0/1 argument for whether to resume the 30 | optimizer and LR schedule, or just the model itself 31 | - ``--mixed-precision`` is a 0/1 argument for whether to use mixed-precision 32 | training or not (required PyTorch compiled with AMP support) 33 | - Model and data loading: 34 | - DataParallel is now *off* by default when loading models, even when 35 | resume_path is specified (previously it was off for new models, and on 36 | for resumed models by default) 37 | - New ``add_custom_forward`` for ``make_and_restore_model`` (see docs for 38 | more details) 39 | - Can now pass a random seed for training data subsetting 40 | - Training: 41 | - See new CLI features---most have training-as-a-library counterparts 42 | - Fixed a bug that did not resume the optimizer and schedule 43 | - Support for custom accuracy functions 44 | - Can now disable ``torch.nograd`` for test set eval (in case you have a 45 | custom accuracy function that needs gradients even on the val set) 46 | - PGD: 47 | - Better random start for l2 attacks 48 | - Added a ``RandomStep`` attacker step (useful for large-noise training with 49 | varying noise over training) 50 | - Fixed bug in the ``with_image`` argument (minor) 51 | - Model saving: 52 | - Accuracies are now saved in the checkpoint files themselves (instead of 53 | just in the log stores) 54 | - Removed redundant checkpoints table from the log store, as it is a 55 | duplicate of the latest checkpoint file and just wastes space 56 | - Cleanup: 57 | - Remove redundant ``save_checkpoint`` function in helpers file 58 | - Code flow improvements 59 | 60 | 61 | robustness 1.1.post2 62 | ''''''''''''''''''''' 63 | - Critical fix in :meth:`robustness.loaders.TransformedLoader`, allow for data shuffling 64 | 65 | robustness 1.1 66 | '''''''''''''' 67 | - Added ability to superclass ImageNet to make 68 | custom datasets (:doc:`docs `) 69 | - Added ``shuffle_train`` and ``shuffle_test`` options to 70 | :meth:`~robustness.datasets.Dataset.make_loaders` 71 | - Added support for cyclic learning rate (``--custom-schedule cyclic`` via command line or ``{"custom_schedule": "cyclic"}`` from Python 72 | - Added support for transfer learning/partial parameter updates, 73 | :meth:`robustness.train.train_model` now takes ``update_params`` argument, 74 | list of parameters to update 75 | - Allow ``random_start`` (random start for adversarial attacks) to be set via 76 | command line 77 | - Change defaults for ImageNet training (``200`` epochs instead of ``350``) 78 | - Small fixes/refinements to :mod:`robustness.tools.vis_tools` module 79 | -------------------------------------------------------------------------------- /docs/example_usage/cli_usage.rst: -------------------------------------------------------------------------------- 1 | Training and evaluating networks via command line 2 | ================================================== 3 | In this walkthrough, we'll go over how to train and evaluate networks via the 4 | :mod:`robustness.main` command-line tool. 5 | 6 | Training a standard (nonrobust) model 7 | -------------------------------------- 8 | We'll start by training a standard (non-robust) model. This is accomplished through the following command: 9 | 10 | .. code-block:: bash 11 | 12 | python -m robustness.main --dataset DATASET --data /path/to/dataset \ 13 | --adv-train 0 --arch ARCH --out-dir /logs/checkpoints/dir/ 14 | 15 | In the above, :samp:`DATASET` can be any supported dataset (i.e. in 16 | :attr:`robustness.datasets.DATASETS`). For a demonstration of how to add a 17 | supported dataset, see :ref:`here `. 18 | 19 | With the above command, you should start seeing progress bars indicating that 20 | the training has begun! Note that there are a whole host of arguments that you 21 | can customize in training, including optimizer parameters (e.g. :samp:`--lr`, 22 | :samp:`--weight-decay`, :samp:`--momentum`), logging parameters (e.g. 23 | :samp:`--log-iters`, :samp:`--save-ckpt-iters`), and learning rate schedule. To 24 | see more about these arguments, we run: 25 | 26 | .. code-block:: bash 27 | 28 | python -m robustness --help 29 | 30 | For completeness, the full list of parameters related to *non-robust* training 31 | are below: 32 | 33 | .. code-block:: bash 34 | 35 | --out-dir OUT_DIR where to save training logs and checkpoints (default: 36 | required) 37 | config path for loading in parameters (default: None) 38 | --exp-name EXP_NAME where to save in (inside out_dir) (default: None) 39 | --dataset {imagenet,restricted_imagenet,cifar,cinic,a2b} 40 | (choices: {arg_type}, default: required) 41 | --data DATA path to the dataset (default: /tmp/) 42 | --arch ARCH architecture (see {cifar,imagenet}_models/ (default: 43 | required) 44 | --batch-size BATCH_SIZE 45 | batch size for data loading (default: by dataset) 46 | --workers WORKERS data loading workers (default: 30) 47 | --resume RESUME path to checkpoint to resume from (default: None) 48 | --data-aug {0,1} whether to use data augmentation (choices: {arg_type}, 49 | default: 1) 50 | --epochs EPOCHS number of epochs to train for (default: by dataset) 51 | --lr LR initial learning rate for training (default: 0.1) 52 | --weight_decay WEIGHT_DECAY 53 | SGD weight decay parameter (default: by dataset) 54 | --momentum MOMENTUM SGD momentum parameter (default: 0.9) 55 | --step-lr STEP_LR number of steps between 10x LR drops (default: by 56 | dataset) 57 | --step-lr-gamma GAMMA multiplier for each LR drop (default: 0.1, i.e., 10x drops) 58 | --custom-lr-multiplier CUSTOM_SCHEDULE 59 | LR sched (format: [(epoch, LR),...]) (default: None) 60 | --lr-interpolation {linear, step} 61 | How to interpolate between learning rates (default: step) 62 | --log-iters LOG_ITERS 63 | how frequently (in epochs) to log (default: 5) 64 | --save-ckpt-iters SAVE_CKPT_ITERS 65 | how frequently (epochs) to save (-1 for bash, only 66 | saves best and last) (default: -1) 67 | --mixed-precision {0, 1} 68 | Whether to use mixed-precision training (needs 69 | to be compiled with NVIDIA AMP support) 70 | 71 | Finally, there is one additional argument, :samp:`--adv-eval {0,1}`, that enables 72 | adversarial evaluation of the non-robust model as it is being trained (i.e. 73 | instead of reporting just standard accuracy every few epochs, we'll also report 74 | robust accuracy if :samp:`--adv-eval 1` is added). However, adding this argument 75 | also necessitates the addition of hyperparameters for adversarial attack, which 76 | we cover in the following section. 77 | 78 | Training a robust model (adversarial training) 79 | -------------------------------------------------- 80 | To train a robust model we proceed in the exact same way as for a standard 81 | model, but with a few changes. First, we change :samp:`--adv-train 0` to 82 | :samp:`--adv-train 1` in the training command. Then, we need to make sure to 83 | supply all the necessary hyperparameters for the attack: 84 | 85 | .. code-block:: bash 86 | 87 | --attack-steps ATTACK_STEPS 88 | number of steps for adversarial attack (default: 7) 89 | --constraint {inf,2,unconstrained} 90 | adv constraint (choices: {arg_type}, default: 91 | required) 92 | --eps EPS adversarial perturbation budget (default: required) 93 | --attack-lr ATTACK_LR 94 | step size for PGD (default: required) 95 | --use-best {0,1} if 1 (0) use best (final) PGD step as example 96 | (choices: {arg_type}, default: 1) 97 | --random-restarts RANDOM_RESTARTS 98 | number of random PGD restarts for eval (default: 0) 99 | --custom-eps-multiplier EPS_SCHEDULE 100 | epsilon multiplier sched (same format as LR schedule) 101 | 102 | 103 | Evaluating trained models 104 | ------------------------- 105 | To evaluate a trained model, we use the ``--eval-only`` flag when calling 106 | :mod:`robustness.main`. To evaluate the model for just standard 107 | (not adversarial) accuracy, only the following arguments are required: 108 | 109 | .. code-block:: bash 110 | 111 | python -m robustness.main --dataset DATASET --data /path/to/dataset \ 112 | --eval-only 1 --out-dir OUT_DIR --arch arch --adv-eval 0 \ 113 | --resume PATH_TO_TRAINED_MODEL_CHECKPOINT 114 | 115 | We can also evaluate adversarial accuracy by changing ``--adv-eval 0`` to 116 | ``--adv-eval 1`` and also adding the arguments from the previous section used 117 | for adversarial attacks. 118 | 119 | Examples 120 | -------- 121 | Training a non-robust ResNet-18 for the CIFAR dataset: 122 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 123 | 124 | .. code-block:: bash 125 | 126 | python -m robustness.main --dataset cifar --data /path/to/cifar \ 127 | --adv-train 0 --arch resnet18 --out-dir /logs/checkpoints/dir/ 128 | 129 | Training a robust ResNet-50 for the Restricted-ImageNet dataset: 130 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 131 | 132 | .. code-block:: bash 133 | 134 | CUDA_VISIBLE_DEVICES=1,2,3,4,5,6 python -m robustness.main --dataset restricted_imagenet --data \ 135 | $IMAGENET_PATH --adv-train 1 --arch resnet50 \ 136 | --out-dir /tmp/logs/checkpoints/dir/ --eps 3.0 --attack-lr 0.5 \ 137 | --attack-steps 7 --constraint 2 138 | 139 | Testing the standard and adversarial accuracy of a trained CIFAR-10 model with 140 | L2 norm constraint of 0.5 and 100 L2-PGD steps: 141 | 142 | .. code-block:: bash 143 | 144 | python -m robustness.main --dataset cifar --eval-only 1 --out-dir /tmp/ \ 145 | --arch resnet50 --adv-eval 1 --constraint 2 --eps 0.5 --attack-lr 0.1 \ 146 | --attack-steps 100 --resume path/to/ckpt/checkpoint.pt.best 147 | 148 | Reading and analyzing training results 149 | -------------------------------------- 150 | 151 | By default, the above command will store all the data generated from the 152 | training process above in a subdirectory inside of 153 | :samp:`/logs/checkpoints/dir/`, the path supplied to the :samp:`--out-dir` 154 | argument. The subdirectory will be named by default via a 36 character, randomly 155 | generated unique identifier, but it can be named manually via the 156 | :samp:`--exp-name` argument. By the end of training, the folder structure will 157 | look something like like: 158 | 159 | .. code-block:: bash 160 | 161 | /logs/checkpoints/dir/a9ffc412-595d-4f8c-8e35-41f000cd35ed 162 | checkpoint.latest.pt 163 | checkpoint.best.pt 164 | store.h5 165 | tensorboard/ 166 | save/ 167 | 168 | This is the file structure of a data store from the 169 | `Cox `_ logging library. It contains all the 170 | tables (stored as Pandas dataframes, in HDF5 format) of data we wrote about the 171 | experiment: 172 | 173 | .. code-block:: python 174 | 175 | >>> from cox import store 176 | >>> s = store.Store('/logs/checkpoints/dir/', '6aeae7de-3549-49d5-adb6-52fe04689b4e') 177 | >>> s.tables 178 | {'ckpts': , 'logs': , 'metadata': } 179 | 180 | We can get the metadata by looking at the metadata table and extracting values 181 | we want. For example, if we wanted to get the learning rate, 0.1: 182 | 183 | .. code-block:: python 184 | 185 | >>> s['metadata'].df['lr'] 186 | 0 0.1 187 | Name: lr, dtype: float64 188 | 189 | Or, if we wanted to find out which epoch had the highest validation accuracy: 190 | 191 | .. code-block:: python 192 | 193 | >>> l_df = s['logs'] 194 | >>> ldf[ldf['nat_prec1'] == max(ldf['nat_prec1'].tolist())]['epoch'].tolist()[0] 195 | 32 196 | 197 | In a similar manner, the 'ckpts' table contains all the previous checkpoints, 198 | and the 'logs' table contains logging information pertaining to the training. 199 | Cox allows us to really easily aggregate training logs across different training 200 | runs and compare/analyze them---we recommend taking a look at the `Cox documentation 201 | `_ for more information on how to use it. 202 | 203 | Note that when training models programmatically (as in our walkthrough 204 | :doc:`Part 1 <../example_usage/training_lib_part_1>` and :doc:`Part 2 <../example_usage/training_lib_part_2>`), it is possible to add on custom 205 | logging functionalities and keep track of essentially anything during training. 206 | -------------------------------------------------------------------------------- /docs/example_usage/custom_imagenet.rst: -------------------------------------------------------------------------------- 1 | Creating a custom dataset by superclassing ImageNet 2 | ==================================================== 3 | Often in both adversarial robustness research and otherwise, datasets with the 4 | richness of ImageNet are desired, but without the added complexity of the 1000-way 5 | ILSVRC classification task. A common workaround is to "superclass" ImageNet, 6 | that is, to define a new dataset that contains broad classes which each subsume 7 | several of the original ImageNet classes. 8 | 9 | In this document, we will discuss how to (a) load pre-packaged ImageNet-based 10 | datasets that we've created, and (b) create new custom N-class subset of 11 | ImageNet data by leveraging the WordNet hierarchy to build superclasses. The 12 | robustness library provides functionality to do this via the 13 | :class:`~robustness.datasets.CustomImageNet` and 14 | :class:`~robustness.tools.imagenet_helpers.ImageNetHierarchy` classes. In this 15 | walkthrough, we'll see how to use these classes to browse and use the WordNet 16 | hierarchy to create custom ImageNet-based datasets. 17 | 18 | .. raw:: html 19 | 20 |    Download 22 | a Jupyter notebook containing all the code from this walkthrough!
23 |
24 | 25 | Requirements/Setup 26 | '''''''''''''''''' 27 | To create custom ImageNet datasets, we need (a) the ImageNet dataset to be 28 | downloaded and available in PyTorch-readable format, and (b) the files 29 | ``wordnet.is_a.txt``, ``words.txt`` and ``imagenet_class_index.json``, all 30 | contained within the same directory (all of these files can be obtained from 31 | `the ImageNet website `_. 32 | 33 | Basic Usage: Loading Pre-Packaged ImageNet-based Datasets 34 | '''''''''''''''''''''''''''''''''''''''''''''''''''''''''' 35 | To make things as easy as possible, we've compiled a list of large, but less 36 | complex ImageNet-based datasets. These datasets can be loaded in their 37 | unbalanced or balanced forms, where in the latter we truncate each class to have 38 | the same number of images as the smallest class. We enumerate these datasets 39 | below: 40 | 41 | =============== ================================================== 42 | Dataset Name Classes 43 | =============== ================================================== 44 | living_9 | Dog (n02084071), Bird (n01503061), 45 | | Arthropod (n01767661), Reptile (n01661091), 46 | | Primate (n02469914), Fish (n02512053), 47 | | Feline (n02120997), Bovid (n02401031), 48 | | Amphibian (n01627424) 49 | mixed_10 | Dog (n02084071), Bird (n01503061), 50 | | Insect (n02159955), Monkey (n02484322), 51 | | Car (n02958343), Cat (n02120997), 52 | | Truck (n04490091), Fruit (n13134947), 53 | | Fungus (n12992868), Boat (n02858304) 54 | mixed_13 | Dog (n02084071), Bird (n01503061), 55 | | Insect (n02159955), Furniture (n03405725), 56 | | Fish (n02512053), Monkey (n02484322), 57 | | Car (n02958343), Cat (n02120997), 58 | | Truck (n04490091), Fruit (n13134947), 59 | | Fungus (n12992868), Boat (n02858304), 60 | | Computer (n03082979) 61 | geirhos_16 | Aircraft (n02686568), Bear (n02131653), 62 | | Bicycle (n02834778), Bird (n01503061), 63 | | Boat (n02858304), Bottle (n02876657), 64 | | Car (n02958343), Cat (n02121808), 65 | | Char (n03001627), Clock (n03046257), 66 | | Dog (n02084071), Elephant (n02503517), 67 | | Keyboard (n03614532), Knife (n03623556), 68 | | Oven (n03862676), Truck (n04490091), 69 | big_12 | Dog (n02084071), Structure(n04341686), 70 | | Bird (n01503061), Clothing (n03051540), 71 | | Vehicle(n04576211), Reptile (n01661091), 72 | | Carnivore (n02075296), Insect (n02159955), 73 | | Instrument (n03800933), Food (n07555863), 74 | | Furniture (n03405725), Primate (n02469914), 75 | =============== ================================================== 76 | 77 | Loading any of these datasets (for example, ``mixed_10``) is relatively simple: 78 | 79 | .. code-block:: python 80 | 81 | from robustness import datasets 82 | from robustness.tools.imagenet_helpers import common_superclass_wnid, ImageNetHierarchy 83 | 84 | in_hier = ImageNetHierarchy(in_path, in_info_path) 85 | superclass_wnid = common_superclass_wnid('mixed_10') 86 | class_ranges, label_map = in_hier.get_subclasses(superclass_wnid, balanced=True) 87 | 88 | In the above, :samp:`in_path` should point to a folder with the ImageNet 89 | dataset in ``train`` and ``val`` sub-folders; :samp:`in_info_path` should be the 90 | path to the directory containing the aforementioned files (``wordnet.is_a.txt``, 91 | ``words.txt``, ``imagenet_class_index.json``). 92 | 93 | We can then create a dataset and the corresponding data loader using: 94 | 95 | .. code-block:: python 96 | 97 | custom_dataset = datasets.CustomImageNet(in_path, class_ranges) 98 | train_loader, test_loader = custom_dataset.make_loaders(workers=num_workers, 99 | batch_size=batch_size) 100 | 101 | You're all set! You can then use this :samp:`custom_dataset` and loaders 102 | just as you would any other existing/custom dataset in the robustness 103 | library. For instance, you can visualize training set samples and their 104 | labels using: 105 | 106 | .. code-block:: python 107 | 108 | from robustness.tools.vis_tools import show_image_row 109 | im, lab = next(iter(train_loader)) 110 | show_image_row([im], tlist=[[label_map[int(k)] for k in lab]]) 111 | 112 | Advanced Usage (Making Custom Datasets) Part 1: Browsing the WordNet Hierarchy 113 | ''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''' 114 | The :class:`~robustness.tools.imagenet_helpers.ImageNetHierarchy` class allows 115 | us to probe the WordNet hierarchy and create custom datasets with the desired 116 | number of superclasses. We first create an instance of the 117 | ``ImageNetHierarchy`` class: 118 | 119 | .. code-block:: python 120 | 121 | from robustness.tools.imagenet_helpers import ImageNetHierarchy 122 | in_hier = ImageNetHierarchy(in_path, in_info_path) 123 | 124 | 125 | Again, :samp:`in_path` should point to a folder with the ImageNet 126 | dataset in ``train`` and ``val`` sub-folders; :samp:`in_info_path` should be the 127 | path to the directory containing the aforementioned files (``wordnet.is_a.txt``, 128 | ``words.txt``, ``imagenet_class_index.json``). 129 | 130 | We can now use the :samp:`in_hier` object to probe the ImageNet hierarchy. The 131 | ``wnid_sorted`` attribute, for example, is an iterator over the WordNet IDs, 132 | sorted by the number of descendents they have which are ImageNet classes: 133 | 134 | .. code-block:: python 135 | 136 | for cnt, (wnid, ndesc_in, ndesc_total) in enumerate(in_hier.wnid_sorted): 137 | print(f"WordNet ID: {wnid}, Name: {in_hier.wnid_to_name[wnid]}, #ImageNet descendants: {ndesc_in}") 138 | 139 | Given any WordNet ID, we can also enumerate all of its subclasses of a given 140 | superclass using the ``in_hier.tree`` object and its related methods/attributes: 141 | 142 | .. code-block:: python 143 | 144 | ancestor_wnid = 'n02120997' 145 | print(f"Superclass | WordNet ID: {ancestor_wnid}, Name: {in_hier.wnid_to_name[ancestor_wnid]}") 146 | 147 | for cnt, wnid in enumerate(in_hier.tree['n02120997'].descendants_all): 148 | print(f"Subclass | WordNet ID: {wnid}, Name: {in_hier.wnid_to_name[wnid]}") 149 | 150 | We can filter these subclasses based on whether they correspond to ImageNet 151 | classes using the ``in_wnids`` attribute: 152 | 153 | .. code-block:: python 154 | 155 | ancestor_wnid = 'n02120997' 156 | print(f"Superclass | WordNet ID: {ancestor_wnid}, Name: {in_hier.wnid_to_name[ancestor_wnid]}") 157 | for cnt, wnid in enumerate(in_hier.tree[ancestor_wnid].descendants_all): 158 | if wnid in in_hier.in_wnids: 159 | print(f"ImageNet subclass | WordNet ID: {wnid}, Name: {in_hier.wnid_to_name[wnid]}") 160 | 161 | 162 | Advanced Usage (Making Custom Datasets) Part 2: Making the Datasets 163 | ''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''' 164 | To create a dataset with the desired number of superclasses we use 165 | the :meth:`~robustness.tools.imagenet_helpers.ImageNetHierarchy.get_superclasses` function, 166 | which takes in a desired number of superclasses :samp:`n_classes`, an 167 | (optional) WordNet ID :samp:`ancestor_wnid` that allows us to fix a common 168 | WordNet ancestor for all the classes in our new dataset, and an optional boolean 169 | :samp:`balanced` to get a balanced dataset (where each superclass 170 | has the same number of ImageNet subclasses). 171 | (see :py:meth:`the docstring 172 | ` for 173 | more details). 174 | 175 | .. code-block:: python 176 | 177 | superclass_wnid, class_ranges, label_map = in_hier.get_superclasses(n_classes, 178 | ancestor_wnid=ancestor_wnid, 179 | balanced=balanced) 180 | 181 | This method returns WordNet IDs of chosen superclasses 182 | :samp:`superclass_wnid`, sets of ImageNet subclasses to group together 183 | for each of the superclasses :samp:`class_ranges`, and a mapping from 184 | superclass number to its human-interpretable description :samp:`label_map`. 185 | 186 | You can also directly provide a list of superclass WordNet IDs :samp:`ancestor_wnid` 187 | that you would like to use to build a custom dataset. For instance, some sample superclass 188 | groupings can be found in 189 | py:meth:`~robustness.tools.imagenet_helpers.ImageNetHierarchy.common_superclass_wnid`. 190 | 191 | Once a list of WordNet IDs has been acquired (whether through the method 192 | described here or just manually), we can use the method presented at the 193 | beginning of this article to load the corresponding dataset: 194 | 195 | .. code-block:: python 196 | 197 | custom_dataset = datasets.CustomImageNet(in_path, class_ranges) 198 | train_loader, test_loader = custom_dataset.make_loaders(workers=num_workers, 199 | batch_size=batch_size) 200 | -------------------------------------------------------------------------------- /docs/example_usage/training_lib_part_1.rst: -------------------------------------------------------------------------------- 1 | Using robustness as a general training library (Part 1: Getting started) 2 | ======================================================================== 3 | In the other walkthroughs, we've demonstrated how to use :samp:`robustness` as a 4 | :doc:`command line tool for training and evaluating models `, and how 5 | to use it as a library for :doc:`input manipulation `. 6 | Here, we'll demonstrate how :samp:`robustness` can be used a general library for 7 | experimenting with neural network training. We've found the library has saved us 8 | a tremendous amount of time both writing boilerplate code and making custom 9 | modifications to the training process (one of the primary goals of the library 10 | is to make such modifications simple). 11 | 12 | This walkthrough will be split into two parts: in the first part (this one), 13 | we'll show how to get started with the :samp:`robustness` library, and go 14 | through the process of making a ``main.py`` file for training neural networks. 15 | In the :doc:`second part <../example_usage/training_lib_part_2>`, we'll show how to customize the training 16 | process through custom loss functions, architectures, datasets, logging, and 17 | more. 18 | 19 | .. raw:: html 20 | 21 |   You can follow along using the 23 | source of robustness.main

24 | 25 |    You can also download 27 | a Jupyter notebook containing code from this walkthrough and the next!
28 |
29 | 30 | Step 1: Imports 31 | ---------------- 32 | Our goal in this tutorial will be to make a python file that works nearly 33 | identically to the robustness :doc:`Command-line tool 34 | <../example_usage/cli_usage>`. That is, a user 35 | will be able to call ``python main.py [--arg value ...]`` to train a standard or 36 | robust model. We'll start by importing the necessary modules from the package: 37 | 38 | .. code-block:: python 39 | 40 | from robustness.datasets import DATASETS 41 | from robustness.model_utils import make_and_restore_model 42 | from robustness.train import train_model 43 | from robustness.defaults import check_and_fill_args 44 | from robustness.tools import constants, helpers 45 | from robustness import defaults 46 | 47 | To make life easier, we use `cox `_ (a super 48 | lightweight python logging library) for logging: 49 | 50 | .. code-block:: python 51 | 52 | from cox import utils 53 | from cox import store 54 | 55 | Finally, we'll also need the following external imports: 56 | 57 | .. code-block:: python 58 | 59 | import torch as ch 60 | from argparse import ArgumentParser 61 | import os 62 | 63 | Step 2: Dealing with arguments 64 | ------------------------------- 65 | In this first step, we'll set up an ``args`` object which has all the parameters 66 | we need to train our model. In Step 2.1 we'll show how to use ``argparse`` to 67 | accept user input for specifying parameters via command line; in Step 2.2 we 68 | show how to sanity-check the ``args`` object and fill in reasonable defaults. 69 | 70 | Step 2.1: Setting up command-line args 71 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 72 | The first real step in making our main file is setting up an 73 | ``argparse.ArgumentParser`` object to parse command-line arguments for us. (If 74 | you are not familiar with the python `argparses 75 | `_ module, we recommend looking 76 | there first). Note that if you're not interested in accepting command-line input 77 | for arguments via argparse, you can skip to Step 2.2. 78 | 79 | The ``robustness`` package provides the :mod:`robustness.defaults` module to 80 | make dealing with arguments less painful. In particular, the properties 81 | :attr:`robustness.defaults.TRAINING_ARGS`, :attr:`robustness.defaults.PGD_ARGS`, 82 | and :attr:`robustness.defaults.MODEL_LOADER_ARGS`, contain all of the arguments 83 | (along with default values) needed for training models: 84 | 85 | - :attr:`~robustness.defaults.TRAINING_ARGS` has all of the model training 86 | arguments, like learning rate, momentum, weight decay, learning rate schedule, 87 | etc. 88 | - :attr:`~robustness.defaults.PGD_ARGS` has all of the arguments needed only for 89 | adversarial training, like number of PGD steps, perturbation budget, type of 90 | norm constraint, etc. 91 | - :attr:`~robustness.defaults.MODEL_LOADER_ARGS` has all of the arguments for 92 | instantiating the model and data loaders: dataset, path to dataset, batch 93 | size, number of workers, etc. 94 | 95 | You can take a look at the documentation of :mod:`robustness.defaults` to 96 | learn more about how these attributes are set up and see exactly which arguments 97 | they contain and with what defaults, as well as which arguments are required. The important thing is that the 98 | ``robustness`` package provides the function 99 | :meth:`robustness.defaults.add_args_to_parser` which takes in an arguments 100 | object like the above, and an ``argparse`` parser, and adds the arguments to the 101 | parser: 102 | 103 | .. code-block:: python 104 | 105 | parser = ArgumentParser() 106 | parser = defaults.add_args_to_parser(defaults.MODEL_LOADER_ARGS, parser) 107 | parser = defaults.add_args_to_parser(defaults.TRAINING_ARGS, parser) 108 | parser = defaults.add_args_to_parser(defaults.PGD_ARGS, parser) 109 | # Note that we can add whatever extra arguments we want to the parser here 110 | args = parser.parse_args() 111 | 112 | **Important note:** Even though the arguments objects do specify defaults for 113 | the arguments, these defaults are **not** given to the parser directly. More on 114 | this in Step 2.2. 115 | 116 | If you don't want to use ``argparse`` and already know the values you want to 117 | use for the parameters, you can look at the :mod:`robustness.defaults` 118 | documentation, and hard-code the desired arguments as follows: 119 | 120 | .. code-block:: python 121 | 122 | # Hard-coded base parameters 123 | train_kwargs = { 124 | 'out_dir': "train_out", 125 | 'adv_train': 1, 126 | 'constraint': '2', 127 | 'eps': 0.5, 128 | 'attack_lr': 1.5, 129 | 'attack_steps': 20 130 | } 131 | 132 | # utils.Parameters is just an object wrapper for dicts implementing 133 | # getattr and settattr 134 | train_args = utils.Parameters(train_kwargs) 135 | 136 | Step 2.2: Sanity checks and defaults 137 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 138 | We generally found the ``ArgumentParser`` defaults to be too restrictive for our 139 | applications, and we also wanted to be able to fill in argument defaults even 140 | when we were not using ``ArgumentParser``. Thus, we fill in default arguments 141 | separately via the :meth:`robustness.defaults.check_and_fill_args` function. 142 | This function takes in the ``args`` Namespace object (basically anything 143 | exposing ``setattr`` and ``getattr`` functions), the same ``ARGS`` attributes 144 | discussed above, and a dataset class (used for filling in per-dataset defaults). 145 | The function fills in the defaults it has, and then throws an error if a 146 | required argument is missing: 147 | 148 | .. code-block:: python 149 | 150 | assert args.dataset is not None, "Must provide a dataset" 151 | ds_class = DATASETS[args.dataset] 152 | 153 | args = check_and_fill_args(args, defaults.TRAINING_ARGS, ds_class) 154 | if args.adv_train or args.adv_eval: 155 | args = check_and_fill_args(args, defaults.PGD_ARGS, ds_class) 156 | args = check_and_fill_args(args, defaults.MODEL_LOADER_ARGS, ds_class) 157 | 158 | Note that the :meth:`~robustness.defaults.check_and_fill_args` function is 159 | totally independent of ``argparse``, and can be used even when you don't want to 160 | support command-line arguments. It's a really useful way of sanity checking the 161 | ``args`` object to make sure that there aren't any missing or misspecified arguments. 162 | 163 | Step 3: Creating the model, dataset, and loader 164 | ------------------------------------------------ 165 | The next step is to create the model, dataset, and data loader. We start by 166 | creating the dataset and loaders as follows: 167 | 168 | .. code-block:: python 169 | 170 | # Load up the dataset 171 | data_path = os.path.expandvars(args.data) 172 | dataset = DATASETS[args.dataset](data_path) 173 | 174 | # Make the data loaders 175 | train_loader, val_loader = dataset.make_loaders(args.workers, 176 | args.batch_size, data_aug=bool(args.data_aug)) 177 | 178 | # Prefetches data to improve performance 179 | train_loader = helpers.DataPrefetcher(train_loader) 180 | val_loader = helpers.DataPrefetcher(val_loader) 181 | 182 | We can now create the model by using the 183 | :meth:`robustness.model_utils.make_and_restore_model` function. This function is 184 | used for both creating new models, or (if a ``resume_path`` if passed) loading 185 | previously saved models. 186 | 187 | .. code-block:: python 188 | 189 | model, _ = make_and_restore_model(arch=args.arch, dataset=dataset) 190 | 191 | Step 4: Training the model 192 | --------------------------- 193 | Finally, we create a `cox Store `_ for saving the results of the 194 | training, and call :meth:`robustness.train.train_model` to begin training: 195 | 196 | .. code-block:: python 197 | 198 | # Create the cox store, and save the arguments in a table 199 | store = store.Store(args.out_dir, args.exp_name) 200 | args_dict = args.as_dict() if isinstance(args, utils.Parameters) else vars(args) 201 | schema = store.schema_from_dict(args_dict) 202 | store.add_table('metadata', schema) 203 | store['metadata'].append_row(args_dict) 204 | 205 | model = train_model(args, model, (train_loader, val_loader), store=store) 206 | 207 | -------------------------------------------------------------------------------- /docs/index.rst: -------------------------------------------------------------------------------- 1 | robustness package 2 | ================== 3 | 4 | .. raw:: html 5 | 6 | View on GitHub

8 | 9 | Install via ``pip``: ``pip install robustness`` 10 | 11 | :samp:`robustness` is a package we (students in the `MadryLab `_) created 12 | to make training, evaluating, and exploring neural networks flexible and easy. 13 | We use it in almost all of our projects (whether they involve 14 | adversarial training or not!) and it will be a dependency in many of our 15 | upcoming code releases. A few projects using the library include: 16 | 17 | - `Code `_ for 18 | "Learning Perceptually-Aligned Representations via Adversarial Robustness" 19 | [EIS+19]_ 20 | - `Code `_ for 21 | "Image Synthesis with a Single (Robust) Classifier" [STE+19]_ 22 | - `Code `_ for 23 | "Do Adversarially Robust ImageNet Models Transfer Better?" [SIE+20]_ 24 | - `Code `_ for 25 | "BREEDS: Benchmarks for Subpopulation Shift" [STM20]_ 26 | - `Code `_ for 27 | "Unadversarial Examples: Designing Objects for Robust Vision" [SIE+21]_ 28 | - `Code `_ for 29 | "Certified Patch Robustness via Smoothed Vision Transformers" [SJW+21]_ 30 | 31 | We demonstrate how to use the library in a set of walkthroughs and our API 32 | reference. Functionality provided by the library includes: 33 | 34 | - Training and evaluating standard and robust models for a variety of 35 | datasets/architectures using a :doc:`CLI interface 36 | `. The library also provides support for adding 37 | :ref:`custom datasets ` and :ref:`model architectures 38 | `. 39 | 40 | .. code-block:: bash 41 | 42 | python -m robustness.main --dataset cifar --data /path/to/cifar \ 43 | --adv-train 0 --arch resnet18 --out-dir /logs/checkpoints/dir/ 44 | 45 | - Performing :doc:`input manipulation 46 | ` using robust (or standard) 47 | models---this includes making adversarial examples, inverting representations, 48 | feature visualization, etc. The library offers a variety of optimization 49 | options (e.g. choice between real/estimated gradients, Fourier/pixel basis, 50 | custom loss functions etc.), and is easily extendable. 51 | 52 | .. code-block:: python 53 | 54 | import torch as ch 55 | from robustness.datasets import CIFAR 56 | from robustness.model_utils import make_and_restore_model 57 | 58 | ds = CIFAR('/path/to/cifar') 59 | model, _ = make_and_restore_model(arch='resnet50', dataset=ds, 60 | resume_path='/path/to/model', state_dict_path='model') 61 | model.eval() 62 | attack_kwargs = { 63 | 'constraint': 'inf', # L-inf PGD 64 | 'eps': 0.05, # Epsilon constraint (L-inf norm) 65 | 'step_size': 0.01, # Learning rate for PGD 66 | 'iterations': 100, # Number of PGD steps 67 | 'targeted': True # Targeted attack 68 | 'custom_loss': None # Use default cross-entropy loss 69 | } 70 | 71 | _, test_loader = ds.make_loaders(workers=0, batch_size=10) 72 | im, label = next(iter(test_loader)) 73 | target_label = (label + ch.randint_like(label, high=9)) % 10 74 | adv_out, adv_im = model(im, target_label, make_adv, **attack_kwargs) 75 | 76 | - Importing :samp:`robustness` as a package, which allows for easy training of 77 | neural networks with support for custom loss functions, logging, data loading, 78 | and more! A good introduction can be found in our two-part walkthrough 79 | (:doc:`Part 1 `, :doc:`Part 2 80 | `). 81 | 82 | .. code-block:: python 83 | 84 | from robustness import model_utils, datasets, train, defaults 85 | from robustness.datasets import CIFAR 86 | 87 | # We use cox (http://github.com/MadryLab/cox) to log, store and analyze 88 | # results. Read more at https//cox.readthedocs.io. 89 | from cox.utils import Parameters 90 | import cox.store 91 | 92 | # Hard-coded dataset, architecture, batch size, workers 93 | ds = CIFAR('/path/to/cifar') 94 | m, _ = model_utils.make_and_restore_model(arch='resnet50', dataset=ds) 95 | train_loader, val_loader = ds.make_loaders(batch_size=128, workers=8) 96 | 97 | # Create a cox store for logging 98 | out_store = cox.store.Store(OUT_DIR) 99 | 100 | # Hard-coded base parameters 101 | train_kwargs = { 102 | 'out_dir': "train_out", 103 | 'adv_train': 1, 104 | 'constraint': '2', 105 | 'eps': 0.5, 106 | 'attack_lr': 1.5, 107 | 'attack_steps': 20 108 | } 109 | train_args = Parameters(train_kwargs) 110 | 111 | # Fill whatever parameters are missing from the defaults 112 | train_args = defaults.check_and_fill_args(train_args, 113 | defaults.TRAINING_ARGS, CIFAR) 114 | train_args = defaults.check_and_fill_args(train_args, 115 | defaults.PGD_ARGS, CIFAR) 116 | 117 | # Train a model 118 | train.train_model(train_args, m, (train_loader, val_loader), store=out_store) 119 | 120 | Citation 121 | -------- 122 | If you use this library in your research, cite it as 123 | follows: 124 | 125 | .. code-block:: bibtex 126 | 127 | @misc{robustness, 128 | title={Robustness (Python Library)}, 129 | author={Logan Engstrom and Andrew Ilyas and Hadi Salman and Shibani Santurkar and Dimitris Tsipras}, 130 | year={2019}, 131 | url={https://github.com/MadryLab/robustness} 132 | } 133 | 134 | *(Have you used the package and found it useful? Let us know!)*. 135 | 136 | Walkthroughs 137 | ------------ 138 | 139 | .. toctree:: 140 | example_usage/cli_usage 141 | example_usage/input_space_manipulation 142 | example_usage/training_lib_part_1 143 | example_usage/training_lib_part_2 144 | example_usage/custom_imagenet 145 | example_usage/breeds_datasets 146 | example_usage/changelog 147 | 148 | API Reference 149 | ------------- 150 | 151 | We provide an API reference where we discuss the role of each module and 152 | provide extensive documentation. 153 | 154 | .. toctree:: 155 | api 156 | 157 | 158 | Contributors 159 | ------------- 160 | - `Andrew Ilyas `_ 161 | - `Logan Engstrom `_ 162 | - `Shibani Santurkar `_ 163 | - `Dimitris Tsipras `_ 164 | - `Hadi Salman `_ 165 | 166 | .. [EIS+19] Engstrom L., Ilyas A., Santurkar S., Tsipras D., Tran B., Madry A. (2019). Learning Perceptually-Aligned Representations via Adversarial Robustness. arXiv, arXiv:1906.00945 167 | 168 | .. [STE+19] Santurkar S., Tsipras D., Tran B., Ilyas A., Engstrom L., Madry A. (2019). Image Synthesis with a Single (Robust) Classifier. arXiv, arXiv:1906.09453 169 | 170 | .. [SIE+20] Salman H., Ilyas A., Engstrom L., Kapoor A., Madry A. (2020). Do Adversarially Robust ImageNet Models Transfer Better? arXiv, arXiv:2007.08489 171 | 172 | .. [STM20] Santurkar S., Tsipras D., Madry A. (2020). BREEDS: Benchmarks for Subpopulation Shift. arXiv, arXiv:2008.04859 173 | 174 | .. [SIE+21] Salman H., Ilyas A., Engstrom L., Vemprala S., Kapoor A., Madry A. (2021). Unadversarial Examples: Designing Objects for Robust Vision. arXiv, arXiv:2012.12235 175 | 176 | .. [SJW+21] Salman H., Jain S., Wong E., Madry A. (2021). : Certified Patch Robustness via Smoothed Vision Transformers. arXiv, arXiv:2110.07719 177 | -------------------------------------------------------------------------------- /docs/make.bat: -------------------------------------------------------------------------------- 1 | @ECHO OFF 2 | 3 | pushd %~dp0 4 | 5 | REM Command file for Sphinx documentation 6 | 7 | if "%SPHINXBUILD%" == "" ( 8 | set SPHINXBUILD=sphinx-build 9 | ) 10 | set SOURCEDIR=. 11 | set BUILDDIR=_build 12 | 13 | if "%1" == "" goto help 14 | 15 | %SPHINXBUILD% >NUL 2>NUL 16 | if errorlevel 9009 ( 17 | echo. 18 | echo.The 'sphinx-build' command was not found. Make sure you have Sphinx 19 | echo.installed, then set the SPHINXBUILD environment variable to point 20 | echo.to the full path of the 'sphinx-build' executable. Alternatively you 21 | echo.may add the Sphinx directory to PATH. 22 | echo. 23 | echo.If you don't have Sphinx installed, grab it from 24 | echo.http://sphinx-doc.org/ 25 | exit /b 1 26 | ) 27 | 28 | %SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% 29 | goto end 30 | 31 | :help 32 | %SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% 33 | 34 | :end 35 | popd 36 | -------------------------------------------------------------------------------- /docs/requirements.txt: -------------------------------------------------------------------------------- 1 | sphinx-automodapi 2 | sphinx-rtd-theme 3 | sphinxcontrib-applehelp 4 | sphinxcontrib-devhelp 5 | sphinxcontrib-fulltoc 6 | sphinxcontrib-htmlhelp 7 | sphinxcontrib-jsmath 8 | sphinxcontrib-qthelp 9 | sphinxcontrib-serializinghtml 10 | sphinx_fontawesome 11 | -------------------------------------------------------------------------------- /notebooks/Input manipulation with pre-trained models.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "ATTACK_EPS = 0.5\n", 10 | "ATTACK_STEPSIZE = 0.1\n", 11 | "ATTACK_STEPS = 10\n", 12 | "NUM_WORKERS = 8\n", 13 | "BATCH_SIZE = 10" 14 | ] 15 | }, 16 | { 17 | "cell_type": "markdown", 18 | "metadata": {}, 19 | "source": [ 20 | "## Set up dataset" 21 | ] 22 | }, 23 | { 24 | "cell_type": "code", 25 | "execution_count": null, 26 | "metadata": {}, 27 | "outputs": [], 28 | "source": [ 29 | "import torch as ch\n", 30 | "from robustness.datasets import CIFAR\n", 31 | "ds = CIFAR('/tmp')" 32 | ] 33 | }, 34 | { 35 | "cell_type": "markdown", 36 | "metadata": {}, 37 | "source": [ 38 | "## Instantiate model" 39 | ] 40 | }, 41 | { 42 | "cell_type": "code", 43 | "execution_count": null, 44 | "metadata": {}, 45 | "outputs": [], 46 | "source": [ 47 | "from robustness.model_utils import make_and_restore_model\n", 48 | "model, _ = make_and_restore_model(arch='resnet50', dataset=ds,\n", 49 | " resume_path='/data/theory/robustopt/robust_models/cifar_l2_eps_05/checkpoint.pt.best')\n", 50 | "model.eval()\n", 51 | "pass" 52 | ] 53 | }, 54 | { 55 | "cell_type": "markdown", 56 | "metadata": {}, 57 | "source": [ 58 | "## Set up loaders" 59 | ] 60 | }, 61 | { 62 | "cell_type": "code", 63 | "execution_count": null, 64 | "metadata": {}, 65 | "outputs": [], 66 | "source": [ 67 | "_, test_loader = ds.make_loaders(workers=NUM_WORKERS, batch_size=BATCH_SIZE)\n", 68 | "_, (im, label) = next(enumerate(test_loader))" 69 | ] 70 | }, 71 | { 72 | "cell_type": "markdown", 73 | "metadata": {}, 74 | "source": [ 75 | "## Generating untargeted adversarial examples" 76 | ] 77 | }, 78 | { 79 | "cell_type": "code", 80 | "execution_count": null, 81 | "metadata": {}, 82 | "outputs": [], 83 | "source": [ 84 | "kwargs = {\n", 85 | " 'constraint':'2', # use L2-PGD\n", 86 | " 'eps': ATTACK_EPS, # L2 radius around original image\n", 87 | " 'step_size': ATTACK_STEPSIZE,\n", 88 | " 'iterations': ATTACK_STEPS,\n", 89 | " 'do_tqdm': True,\n", 90 | "}" 91 | ] 92 | }, 93 | { 94 | "cell_type": "code", 95 | "execution_count": null, 96 | "metadata": {}, 97 | "outputs": [], 98 | "source": [ 99 | "_, im_adv = model(im, label, make_adv=True, **kwargs)" 100 | ] 101 | }, 102 | { 103 | "cell_type": "code", 104 | "execution_count": null, 105 | "metadata": {}, 106 | "outputs": [], 107 | "source": [ 108 | "from robustness.tools.vis_tools import show_image_row\n", 109 | "from robustness.tools.label_maps import CLASS_DICT\n", 110 | "\n", 111 | "# Get predicted labels for adversarial examples\n", 112 | "pred, _ = model(im_adv)\n", 113 | "label_pred = ch.argmax(pred, dim=1)\n", 114 | "\n", 115 | "# Visualize test set images, along with corresponding adversarial examples\n", 116 | "show_image_row([im.cpu(), im_adv.cpu()],\n", 117 | " tlist=[[CLASS_DICT['CIFAR'][int(t)] for t in l] for l in [label, label_pred]],\n", 118 | " fontsize=18,\n", 119 | " filename='./adversarial_example_CIFAR.png')" 120 | ] 121 | }, 122 | { 123 | "cell_type": "markdown", 124 | "metadata": {}, 125 | "source": [ 126 | "## Targeted adversarial examples" 127 | ] 128 | }, 129 | { 130 | "cell_type": "code", 131 | "execution_count": null, 132 | "metadata": {}, 133 | "outputs": [], 134 | "source": [ 135 | "kwargs = {\n", 136 | " 'constraint':'2',\n", 137 | " 'eps': ATTACK_EPS,\n", 138 | " 'step_size': ATTACK_STEPSIZE,\n", 139 | " 'iterations': ATTACK_STEPS,\n", 140 | " 'targeted': True,\n", 141 | " 'do_tqdm': True\n", 142 | "}" 143 | ] 144 | }, 145 | { 146 | "cell_type": "code", 147 | "execution_count": null, 148 | "metadata": {}, 149 | "outputs": [], 150 | "source": [ 151 | "targ = ch.zeros_like(label)" 152 | ] 153 | }, 154 | { 155 | "cell_type": "code", 156 | "execution_count": null, 157 | "metadata": {}, 158 | "outputs": [], 159 | "source": [ 160 | "_, im_adv = model(im, targ, make_adv=True, **kwargs)" 161 | ] 162 | }, 163 | { 164 | "cell_type": "code", 165 | "execution_count": null, 166 | "metadata": {}, 167 | "outputs": [], 168 | "source": [ 169 | "# Visualize test set images, along with corresponding adversarial examples\n", 170 | "show_image_row([im.cpu(), im_adv.cpu()],\n", 171 | " tlist=[[CLASS_DICT['CIFAR'][int(t)] for t in l] for l in [label, label_pred]],\n", 172 | " fontsize=18,\n", 173 | " filename='./adversarial_example_CIFAR.png')" 174 | ] 175 | }, 176 | { 177 | "cell_type": "markdown", 178 | "metadata": {}, 179 | "source": [ 180 | "## Custom Input Manipulation (Representation Inversion)" 181 | ] 182 | }, 183 | { 184 | "cell_type": "code", 185 | "execution_count": null, 186 | "metadata": {}, 187 | "outputs": [], 188 | "source": [ 189 | "_, (im_inv, label_inv) = next(enumerate(test_loader)) # Images to invert\n", 190 | "with ch.no_grad():\n", 191 | " (_, rep_inv), _ = model(im_inv, with_latent=True) # Corresponding representation" 192 | ] 193 | }, 194 | { 195 | "cell_type": "code", 196 | "execution_count": null, 197 | "metadata": {}, 198 | "outputs": [], 199 | "source": [ 200 | "def inversion_loss(model, inp, targ):\n", 201 | " # Compute representation for the input\n", 202 | " _, rep = model(inp, with_latent=True, fake_relu=True)\n", 203 | " # Normalized L2 error w.r.t. the target representation\n", 204 | " loss = ch.div(ch.norm(rep - targ, dim=1), ch.norm(targ, dim=1))\n", 205 | " return loss, None" 206 | ] 207 | }, 208 | { 209 | "cell_type": "code", 210 | "execution_count": null, 211 | "metadata": {}, 212 | "outputs": [], 213 | "source": [ 214 | "kwargs = {\n", 215 | " 'custom_loss': inversion_loss,\n", 216 | " 'constraint':'2',\n", 217 | " 'eps': 1000,\n", 218 | " 'step_size': 1,\n", 219 | " 'iterations': 1000,\n", 220 | " 'targeted': True,\n", 221 | " 'do_tqdm': True,\n", 222 | "}" 223 | ] 224 | }, 225 | { 226 | "cell_type": "code", 227 | "execution_count": null, 228 | "metadata": {}, 229 | "outputs": [], 230 | "source": [ 231 | "im_seed = ch.clamp(ch.randn_like(im_inv) / 20 + 0.5, 0, 1)" 232 | ] 233 | }, 234 | { 235 | "cell_type": "code", 236 | "execution_count": null, 237 | "metadata": {}, 238 | "outputs": [], 239 | "source": [ 240 | "_, im_matched = model(im_seed, rep_inv, make_adv=True, **kwargs)" 241 | ] 242 | }, 243 | { 244 | "cell_type": "code", 245 | "execution_count": null, 246 | "metadata": {}, 247 | "outputs": [], 248 | "source": [ 249 | "show_image_row([im_inv.cpu(), im_seed.cpu(), im_matched.cpu()],\n", 250 | " [\"Original\", r\"Seed ($x_0$)\", \"Result\"],\n", 251 | " fontsize=18,\n", 252 | " filename='./custom_inversion_CIFAR.png')" 253 | ] 254 | }, 255 | { 256 | "cell_type": "markdown", 257 | "metadata": {}, 258 | "source": [ 259 | "## Other optimization methods (example: Fourier basis)" 260 | ] 261 | }, 262 | { 263 | "cell_type": "code", 264 | "execution_count": null, 265 | "metadata": {}, 266 | "outputs": [], 267 | "source": [ 268 | "kwargs = {\n", 269 | " 'custom_loss': inversion_loss,\n", 270 | " 'constraint':'fourier',\n", 271 | " 'eps': 1000, # ignored anyways\n", 272 | " 'step_size': 500, # have to re-tune LR\n", 273 | " 'iterations': 10000,\n", 274 | " 'targeted': True,\n", 275 | " 'do_tqdm': True,\n", 276 | "}" 277 | ] 278 | }, 279 | { 280 | "cell_type": "code", 281 | "execution_count": null, 282 | "metadata": {}, 283 | "outputs": [], 284 | "source": [ 285 | "im_seed = ch.randn(BATCH_SIZE, 3, 32, 32, 2) / 5" 286 | ] 287 | }, 288 | { 289 | "cell_type": "code", 290 | "execution_count": null, 291 | "metadata": {}, 292 | "outputs": [], 293 | "source": [ 294 | "_, im_matched = model(im_seed, rep_inv, make_adv=True, **kwargs)" 295 | ] 296 | }, 297 | { 298 | "cell_type": "code", 299 | "execution_count": null, 300 | "metadata": {}, 301 | "outputs": [], 302 | "source": [ 303 | "show_image_row([im_inv.cpu(), im_matched.cpu()],\n", 304 | " [\"Original\", \"Result\"],\n", 305 | " fontsize=18,\n", 306 | " filename='./custom_inversion_CIFAR_fourier.png')" 307 | ] 308 | } 309 | ], 310 | "metadata": { 311 | "kernelspec": { 312 | "display_name": "Python 3", 313 | "language": "python", 314 | "name": "python3" 315 | }, 316 | "language_info": { 317 | "codemirror_mode": { 318 | "name": "ipython", 319 | "version": 3 320 | }, 321 | "file_extension": ".py", 322 | "mimetype": "text/x-python", 323 | "name": "python", 324 | "nbconvert_exporter": "python", 325 | "pygments_lexer": "ipython3", 326 | "version": "3.6.8" 327 | } 328 | }, 329 | "nbformat": 4, 330 | "nbformat_minor": 2 331 | } 332 | -------------------------------------------------------------------------------- /notebooks/Using robustness as a library.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "OUT_DIR = '/tmp/'\n", 10 | "NUM_WORKERS = 16\n", 11 | "BATCH_SIZE = 512" 12 | ] 13 | }, 14 | { 15 | "cell_type": "markdown", 16 | "metadata": {}, 17 | "source": [ 18 | "# Barebones starter example" 19 | ] 20 | }, 21 | { 22 | "cell_type": "markdown", 23 | "metadata": {}, 24 | "source": [ 25 | "### Imports" 26 | ] 27 | }, 28 | { 29 | "cell_type": "code", 30 | "execution_count": null, 31 | "metadata": {}, 32 | "outputs": [], 33 | "source": [ 34 | "from robustness import model_utils, datasets, train, defaults\n", 35 | "from robustness.datasets import CIFAR\n", 36 | "import torch as ch\n", 37 | "\n", 38 | "# We use cox (http://github.com/MadryLab/cox) to log, store and analyze\n", 39 | "# results. Read more at https//cox.readthedocs.io.\n", 40 | "from cox.utils import Parameters\n", 41 | "import cox.store" 42 | ] 43 | }, 44 | { 45 | "cell_type": "markdown", 46 | "metadata": {}, 47 | "source": [ 48 | "### Make dataset and loaders" 49 | ] 50 | }, 51 | { 52 | "cell_type": "code", 53 | "execution_count": null, 54 | "metadata": {}, 55 | "outputs": [], 56 | "source": [ 57 | "# Hard-coded dataset, architecture, batch size, workers\n", 58 | "ds = CIFAR('/tmp/')\n", 59 | "m, _ = model_utils.make_and_restore_model(arch='resnet18', dataset=ds)\n", 60 | "train_loader, val_loader = ds.make_loaders(batch_size=BATCH_SIZE, workers=NUM_WORKERS)" 61 | ] 62 | }, 63 | { 64 | "cell_type": "markdown", 65 | "metadata": {}, 66 | "source": [ 67 | "### Make a cox store for logging" 68 | ] 69 | }, 70 | { 71 | "cell_type": "code", 72 | "execution_count": null, 73 | "metadata": {}, 74 | "outputs": [], 75 | "source": [ 76 | "# Create a cox store for logging\n", 77 | "out_store = cox.store.Store(OUT_DIR)" 78 | ] 79 | }, 80 | { 81 | "cell_type": "markdown", 82 | "metadata": {}, 83 | "source": [ 84 | "### Set up training arguments" 85 | ] 86 | }, 87 | { 88 | "cell_type": "code", 89 | "execution_count": null, 90 | "metadata": {}, 91 | "outputs": [], 92 | "source": [ 93 | "# Hard-coded base parameters\n", 94 | "train_kwargs = {\n", 95 | " 'out_dir': \"train_out\",\n", 96 | " 'adv_train': 1,\n", 97 | " 'constraint': '2',\n", 98 | " 'eps': 0.5,\n", 99 | " 'attack_lr': 0.1,\n", 100 | " 'attack_steps': 7,\n", 101 | " 'epochs': 5\n", 102 | "}\n", 103 | "train_args = Parameters(train_kwargs)\n", 104 | "\n", 105 | "# Fill whatever parameters are missing from the defaults\n", 106 | "train_args = defaults.check_and_fill_args(train_args,\n", 107 | " defaults.TRAINING_ARGS, CIFAR)\n", 108 | "train_args = defaults.check_and_fill_args(train_args,\n", 109 | " defaults.PGD_ARGS, CIFAR)" 110 | ] 111 | }, 112 | { 113 | "cell_type": "markdown", 114 | "metadata": {}, 115 | "source": [ 116 | "### Train Model" 117 | ] 118 | }, 119 | { 120 | "cell_type": "code", 121 | "execution_count": null, 122 | "metadata": {}, 123 | "outputs": [], 124 | "source": [ 125 | "# Train a model\n", 126 | "train.train_model(train_args, m, (train_loader, val_loader), store=out_store)\n", 127 | "pass" 128 | ] 129 | }, 130 | { 131 | "cell_type": "markdown", 132 | "metadata": {}, 133 | "source": [ 134 | "# Customizations" 135 | ] 136 | }, 137 | { 138 | "cell_type": "markdown", 139 | "metadata": {}, 140 | "source": [ 141 | "## Custom loss" 142 | ] 143 | }, 144 | { 145 | "cell_type": "code", 146 | "execution_count": null, 147 | "metadata": {}, 148 | "outputs": [], 149 | "source": [ 150 | "train_crit = ch.nn.CrossEntropyLoss()\n", 151 | "def custom_train_loss(logits, targ):\n", 152 | " probs = ch.ones_like(logits) * 0.5\n", 153 | " logits_to_multiply = ch.bernoulli(probs) * 9 + 1\n", 154 | " return train_crit(logits_to_multiply * logits, targ)\n", 155 | "\n", 156 | "adv_crit = ch.nn.CrossEntropyLoss(reduction='none').cuda()\n", 157 | "def custom_adv_loss(model, inp, targ):\n", 158 | " logits = model(inp)\n", 159 | " probs = ch.ones_like(logits) * 0.5\n", 160 | " logits_to_multiply = ch.bernoulli(probs) * 9 + 1\n", 161 | " new_logits = logits_to_multiply * logits\n", 162 | " return adv_crit(new_logits, targ), new_logits\n", 163 | "\n", 164 | "train_args.custom_train_loss = custom_train_loss\n", 165 | "train_args.custom_adv_loss = custom_adv_loss" 166 | ] 167 | }, 168 | { 169 | "cell_type": "code", 170 | "execution_count": null, 171 | "metadata": {}, 172 | "outputs": [], 173 | "source": [ 174 | "train.train_model(train_args, m, (train_loader, val_loader), store=out_store)" 175 | ] 176 | }, 177 | { 178 | "cell_type": "markdown", 179 | "metadata": {}, 180 | "source": [ 181 | "## Custom data loaders" 182 | ] 183 | }, 184 | { 185 | "cell_type": "markdown", 186 | "metadata": {}, 187 | "source": [ 188 | "### Using LambdaLoader" 189 | ] 190 | }, 191 | { 192 | "cell_type": "code", 193 | "execution_count": null, 194 | "metadata": {}, 195 | "outputs": [], 196 | "source": [ 197 | "from robustness.loaders import LambdaLoader\n", 198 | "\n", 199 | "def label_noiser(ims, labels):\n", 200 | " label_noise = ch.randint_like(labels, high=9)\n", 201 | " probs = ch.ones_like(label_noise) * 0.1\n", 202 | " labels_to_noise = ch.bernoulli(probs.float()).long()\n", 203 | " new_labels = (labels + label_noise * labels_to_noise) % 10\n", 204 | " return ims, new_labels\n", 205 | "\n", 206 | "train_loader = LambdaLoader(train_loader, label_noiser)" 207 | ] 208 | }, 209 | { 210 | "cell_type": "code", 211 | "execution_count": null, 212 | "metadata": {}, 213 | "outputs": [], 214 | "source": [ 215 | "train.train_model(train_args, m, (train_loader, val_loader), store=out_store)\n", 216 | "pass" 217 | ] 218 | }, 219 | { 220 | "cell_type": "markdown", 221 | "metadata": {}, 222 | "source": [ 223 | "### Using TransformedLoader" 224 | ] 225 | }, 226 | { 227 | "cell_type": "code", 228 | "execution_count": null, 229 | "metadata": {}, 230 | "outputs": [], 231 | "source": [ 232 | "from robustness.loaders import TransformedLoader\n", 233 | "from robustness.data_augmentation import TRAIN_TRANSFORMS_DEFAULT\n", 234 | "\n", 235 | "def make_rand_labels(ims, targs):\n", 236 | " new_targs = ch.randint(0, high=10,size=targs.shape).long()\n", 237 | " return ims, new_targs\n", 238 | "\n", 239 | "train_loader_transformed = TransformedLoader(train_loader,\n", 240 | " make_rand_labels,\n", 241 | " TRAIN_TRANSFORMS_DEFAULT(32),\n", 242 | " workers=8,\n", 243 | " batch_size=BATCH_SIZE,\n", 244 | " do_tqdm=True)" 245 | ] 246 | }, 247 | { 248 | "cell_type": "code", 249 | "execution_count": null, 250 | "metadata": {}, 251 | "outputs": [], 252 | "source": [ 253 | "train.train_model(train_args, m, (train_loader, val_loader), store=out_store)\n", 254 | "pass" 255 | ] 256 | }, 257 | { 258 | "cell_type": "markdown", 259 | "metadata": {}, 260 | "source": [ 261 | "## Custom per-iteration logging" 262 | ] 263 | }, 264 | { 265 | "cell_type": "code", 266 | "execution_count": null, 267 | "metadata": {}, 268 | "outputs": [], 269 | "source": [ 270 | "CUSTOM_SCHEMA = {'iteration': int, 'weight_norm': float }\n", 271 | "out_store.add_table('custom', CUSTOM_SCHEMA)" 272 | ] 273 | }, 274 | { 275 | "cell_type": "code", 276 | "execution_count": null, 277 | "metadata": {}, 278 | "outputs": [], 279 | "source": [ 280 | "from torch.nn.utils import parameters_to_vector as flatten\n", 281 | "\n", 282 | "def log_norm(mod, it, loop_type, inp, targ):\n", 283 | " if loop_type == 'train':\n", 284 | " curr_params = flatten(mod.parameters())\n", 285 | " log_info_custom = { 'iteration': it,\n", 286 | " 'weight_norm': ch.norm(curr_params).detach().cpu().numpy() }\n", 287 | " out_store['custom'].append_row(log_info_custom)\n", 288 | " \n", 289 | "train_args.iteration_hook = log_norm" 290 | ] 291 | }, 292 | { 293 | "cell_type": "code", 294 | "execution_count": null, 295 | "metadata": {}, 296 | "outputs": [], 297 | "source": [ 298 | "train.train_model(train_args, m, (train_loader, val_loader), store=out_store)\n", 299 | "pass" 300 | ] 301 | }, 302 | { 303 | "cell_type": "markdown", 304 | "metadata": {}, 305 | "source": [ 306 | "## Custom architecture" 307 | ] 308 | }, 309 | { 310 | "cell_type": "code", 311 | "execution_count": null, 312 | "metadata": {}, 313 | "outputs": [], 314 | "source": [ 315 | "from torch import nn\n", 316 | "from robustness.model_utils import make_and_restore_model\n", 317 | "\n", 318 | "class MLP(nn.Module):\n", 319 | " # Must implement the num_classes argument\n", 320 | " def __init__(self, num_classes=10):\n", 321 | " super().__init__()\n", 322 | " self.fc1 = nn.Linear(32*32*3, 1000)\n", 323 | " self.relu1 = nn.ReLU()\n", 324 | " self.fc2 = nn.Linear(1000, num_classes)\n", 325 | "\n", 326 | " def forward(self, x, *args, **kwargs):\n", 327 | " out = x.view(x.shape[0], -1)\n", 328 | " out = self.fc1(out)\n", 329 | " out = self.relu1(out)\n", 330 | " return self.fc2(out)\n", 331 | "\n", 332 | "new_model = MLP(num_classes=10)" 333 | ] 334 | }, 335 | { 336 | "cell_type": "code", 337 | "execution_count": null, 338 | "metadata": {}, 339 | "outputs": [], 340 | "source": [ 341 | "new_model, _ = make_and_restore_model(arch=new_model, dataset=ds)" 342 | ] 343 | }, 344 | { 345 | "cell_type": "code", 346 | "execution_count": null, 347 | "metadata": {}, 348 | "outputs": [], 349 | "source": [ 350 | "train.train_model(train_args, new_model, (train_loader, val_loader), store=out_store)\n", 351 | "pass" 352 | ] 353 | } 354 | ], 355 | "metadata": { 356 | "kernelspec": { 357 | "display_name": "Python 3", 358 | "language": "python", 359 | "name": "python3" 360 | }, 361 | "language_info": { 362 | "codemirror_mode": { 363 | "name": "ipython", 364 | "version": 3 365 | }, 366 | "file_extension": ".py", 367 | "mimetype": "text/x-python", 368 | "name": "python", 369 | "nbconvert_exporter": "python", 370 | "pygments_lexer": "ipython3", 371 | "version": "3.6.8" 372 | } 373 | }, 374 | "nbformat": 4, 375 | "nbformat_minor": 2 376 | } 377 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | torch 2 | torchvision 3 | pandas 4 | numpy 5 | scipy 6 | GPUtil 7 | dill 8 | tensorboardX 9 | tables 10 | scikit-learn 11 | seaborn 12 | cox 13 | matplotlib 14 | networkx 15 | -------------------------------------------------------------------------------- /robustness/.gitignore: -------------------------------------------------------------------------------- 1 | *tar 2 | -------------------------------------------------------------------------------- /robustness/__init__.py: -------------------------------------------------------------------------------- 1 | name = "robustness" 2 | __version__ = "1.2.1.post2" 3 | -------------------------------------------------------------------------------- /robustness/attack_steps.py: -------------------------------------------------------------------------------- 1 | """ 2 | **For most use cases, this can just be considered an internal class and 3 | ignored.** 4 | 5 | This module contains the abstract class AttackerStep as well as a few subclasses. 6 | 7 | AttackerStep is a generic way to implement optimizers specifically for use with 8 | :class:`robustness.attacker.AttackerModel`. In general, except for when you want 9 | to :ref:`create a custom optimization method `, you probably do not need to 10 | import or edit this module and can just think of it as internal. 11 | """ 12 | 13 | import torch as ch 14 | 15 | class AttackerStep: 16 | ''' 17 | Generic class for attacker steps, under perturbation constraints 18 | specified by an "origin input" and a perturbation magnitude. 19 | Must implement project, step, and random_perturb 20 | ''' 21 | def __init__(self, orig_input, eps, step_size, use_grad=True): 22 | ''' 23 | Initialize the attacker step with a given perturbation magnitude. 24 | 25 | Args: 26 | eps (float): the perturbation magnitude 27 | orig_input (ch.tensor): the original input 28 | ''' 29 | self.orig_input = orig_input 30 | self.eps = eps 31 | self.step_size = step_size 32 | self.use_grad = use_grad 33 | 34 | def project(self, x): 35 | ''' 36 | Given an input x, project it back into the feasible set 37 | 38 | Args: 39 | ch.tensor x : the input to project back into the feasible set. 40 | 41 | Returns: 42 | A `ch.tensor` that is the input projected back into 43 | the feasible set, that is, 44 | .. math:: \min_{x' \in S} \|x' - x\|_2 45 | ''' 46 | raise NotImplementedError 47 | 48 | def step(self, x, g): 49 | ''' 50 | Given a gradient, make the appropriate step according to the 51 | perturbation constraint (e.g. dual norm maximization for :math:`\ell_p` 52 | norms). 53 | 54 | Parameters: 55 | g (ch.tensor): the raw gradient 56 | 57 | Returns: 58 | The new input, a ch.tensor for the next step. 59 | ''' 60 | raise NotImplementedError 61 | 62 | def random_perturb(self, x): 63 | ''' 64 | Given a starting input, take a random step within the feasible set 65 | ''' 66 | raise NotImplementedError 67 | 68 | def to_image(self, x): 69 | ''' 70 | Given an input (which may be in an alternative parameterization), 71 | convert it to a valid image (this is implemented as the identity 72 | function by default as most of the time we use the pixel 73 | parameterization, but for alternative parameterizations this functino 74 | must be overriden). 75 | ''' 76 | return x 77 | 78 | ### Instantiations of the AttackerStep class 79 | 80 | # L-infinity threat model 81 | class LinfStep(AttackerStep): 82 | """ 83 | Attack step for :math:`\ell_\infty` threat model. Given :math:`x_0` 84 | and :math:`\epsilon`, the constraint set is given by: 85 | 86 | .. math:: S = \{x | \|x - x_0\|_\infty \leq \epsilon\} 87 | """ 88 | def project(self, x): 89 | """ 90 | """ 91 | diff = x - self.orig_input 92 | diff = ch.clamp(diff, -self.eps, self.eps) 93 | return ch.clamp(diff + self.orig_input, 0, 1) 94 | 95 | def step(self, x, g): 96 | """ 97 | """ 98 | step = ch.sign(g) * self.step_size 99 | return x + step 100 | 101 | def random_perturb(self, x): 102 | """ 103 | """ 104 | new_x = x + 2 * (ch.rand_like(x) - 0.5) * self.eps 105 | return ch.clamp(new_x, 0, 1) 106 | 107 | # L2 threat model 108 | class L2Step(AttackerStep): 109 | """ 110 | Attack step for :math:`\ell_\infty` threat model. Given :math:`x_0` 111 | and :math:`\epsilon`, the constraint set is given by: 112 | 113 | .. math:: S = \{x | \|x - x_0\|_2 \leq \epsilon\} 114 | """ 115 | def project(self, x): 116 | """ 117 | """ 118 | diff = x - self.orig_input 119 | diff = diff.renorm(p=2, dim=0, maxnorm=self.eps) 120 | return ch.clamp(self.orig_input + diff, 0, 1) 121 | 122 | def step(self, x, g): 123 | """ 124 | """ 125 | l = len(x.shape) - 1 126 | g_norm = ch.norm(g.view(g.shape[0], -1), dim=1).view(-1, *([1]*l)) 127 | scaled_g = g / (g_norm + 1e-10) 128 | return x + scaled_g * self.step_size 129 | 130 | def random_perturb(self, x): 131 | """ 132 | """ 133 | l = len(x.shape) - 1 134 | rp = ch.randn_like(x) 135 | rp_norm = rp.view(rp.shape[0], -1).norm(dim=1).view(-1, *([1]*l)) 136 | return ch.clamp(x + self.eps * rp / (rp_norm + 1e-10), 0, 1) 137 | 138 | # Unconstrained threat model 139 | class UnconstrainedStep(AttackerStep): 140 | """ 141 | Unconstrained threat model, :math:`S = [0, 1]^n`. 142 | """ 143 | def project(self, x): 144 | """ 145 | """ 146 | return ch.clamp(x, 0, 1) 147 | 148 | def step(self, x, g): 149 | """ 150 | """ 151 | return x + g * self.step_size 152 | 153 | def random_perturb(self, x): 154 | """ 155 | """ 156 | new_x = x + (ch.rand_like(x) - 0.5).renorm(p=2, dim=0, maxnorm=step_size) 157 | return ch.clamp(new_x, 0, 1) 158 | 159 | class FourierStep(AttackerStep): 160 | """ 161 | Step under the Fourier (decorrelated) parameterization of an image. 162 | 163 | See https://distill.pub/2017/feature-visualization/#preconditioning for more information. 164 | """ 165 | def project(self, x): 166 | """ 167 | """ 168 | return x 169 | 170 | def step(self, x, g): 171 | """ 172 | """ 173 | return x + g * self.step_size 174 | 175 | def random_perturb(self, x): 176 | """ 177 | """ 178 | return x 179 | 180 | def to_image(self, x): 181 | """ 182 | """ 183 | return ch.sigmoid(ch.irfft(x, 2, normalized=True, onesided=False)) 184 | 185 | class RandomStep(AttackerStep): 186 | """ 187 | Step for Randomized Smoothing. 188 | """ 189 | def __init__(self, *args, **kwargs): 190 | super().__init__(*args, **kwargs) 191 | self.use_grad = False 192 | 193 | def project(self, x): 194 | """ 195 | """ 196 | return x 197 | 198 | def step(self, x, g): 199 | """ 200 | """ 201 | return x + self.step_size * ch.randn_like(x) 202 | 203 | def random_perturb(self, x): 204 | """ 205 | """ 206 | return x 207 | -------------------------------------------------------------------------------- /robustness/cifar_models/__init__.py: -------------------------------------------------------------------------------- 1 | from .resnet import * 2 | from .vgg import * 3 | from .densenet import * 4 | from .inception import * 5 | -------------------------------------------------------------------------------- /robustness/cifar_models/densenet.py: -------------------------------------------------------------------------------- 1 | '''DenseNet in PyTorch.''' 2 | import math 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | from ..tools.custom_modules import FakeReLU 8 | 9 | 10 | class Bottleneck(nn.Module): 11 | def __init__(self, in_planes, growth_rate): 12 | super(Bottleneck, self).__init__() 13 | self.bn1 = nn.BatchNorm2d(in_planes) 14 | self.conv1 = nn.Conv2d(in_planes, 4*growth_rate, kernel_size=1, bias=False) 15 | self.bn2 = nn.BatchNorm2d(4*growth_rate) 16 | self.conv2 = nn.Conv2d(4*growth_rate, growth_rate, kernel_size=3, padding=1, bias=False) 17 | 18 | def forward(self, x): 19 | out = self.conv1(F.relu(self.bn1(x))) 20 | out = self.conv2(F.relu(self.bn2(out))) 21 | out = torch.cat([out,x], 1) 22 | return out 23 | 24 | 25 | class Transition(nn.Module): 26 | def __init__(self, in_planes, out_planes): 27 | super(Transition, self).__init__() 28 | self.bn = nn.BatchNorm2d(in_planes) 29 | self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=1, bias=False) 30 | 31 | def forward(self, x): 32 | out = self.conv(F.relu(self.bn(x))) 33 | out = F.avg_pool2d(out, 2) 34 | return out 35 | 36 | 37 | class DenseNet(nn.Module): 38 | def __init__(self, block, nblocks, growth_rate=12, reduction=0.5, num_classes=10): 39 | super(DenseNet, self).__init__() 40 | self.growth_rate = growth_rate 41 | 42 | num_planes = 2*growth_rate 43 | self.conv1 = nn.Conv2d(3, num_planes, kernel_size=3, padding=1, bias=False) 44 | 45 | self.dense1 = self._make_dense_layers(block, num_planes, nblocks[0]) 46 | num_planes += nblocks[0]*growth_rate 47 | out_planes = int(math.floor(num_planes*reduction)) 48 | self.trans1 = Transition(num_planes, out_planes) 49 | num_planes = out_planes 50 | 51 | self.dense2 = self._make_dense_layers(block, num_planes, nblocks[1]) 52 | num_planes += nblocks[1]*growth_rate 53 | out_planes = int(math.floor(num_planes*reduction)) 54 | self.trans2 = Transition(num_planes, out_planes) 55 | num_planes = out_planes 56 | 57 | self.dense3 = self._make_dense_layers(block, num_planes, nblocks[2]) 58 | num_planes += nblocks[2]*growth_rate 59 | out_planes = int(math.floor(num_planes*reduction)) 60 | self.trans3 = Transition(num_planes, out_planes) 61 | num_planes = out_planes 62 | 63 | self.dense4 = self._make_dense_layers(block, num_planes, nblocks[3]) 64 | num_planes += nblocks[3]*growth_rate 65 | 66 | self.bn = nn.BatchNorm2d(num_planes) 67 | self.linear = nn.Linear(num_planes, num_classes) 68 | 69 | def _make_dense_layers(self, block, in_planes, nblock): 70 | layers = [] 71 | for i in range(nblock): 72 | layers.append(block(in_planes, self.growth_rate)) 73 | in_planes += self.growth_rate 74 | return nn.Sequential(*layers) 75 | 76 | def forward(self, x, with_latent=False, fake_relu=False, no_relu=False): 77 | assert not no_relu, \ 78 | "DenseNet has no pre-ReLU activations, no_relu not supported" 79 | out = self.conv1(x) 80 | out = self.trans1(self.dense1(out)) 81 | out = self.trans2(self.dense2(out)) 82 | out = self.trans3(self.dense3(out)) 83 | out = self.dense4(out) 84 | if fake_relu: 85 | out = F.avg_pool2d(FakeReLU.apply(self.bn(out)), 4) 86 | else: 87 | out = F.avg_pool2d(F.relu(self.bn(out)), 4) 88 | out = out.view(out.size(0), -1) 89 | latent = out.clone() 90 | out = self.linear(out) 91 | if with_latent: 92 | return out, latent 93 | return out 94 | 95 | def DenseNet121(**kwargs): 96 | return DenseNet(Bottleneck, [6,12,24,16], growth_rate=32, **kwargs) 97 | 98 | def DenseNet169(**kwargs): 99 | return DenseNet(Bottleneck, [6,12,32,32], growth_rate=32, **kwargs) 100 | 101 | def DenseNet201(**kwargs): 102 | return DenseNet(Bottleneck, [6,12,48,32], growth_rate=32, **kwargs) 103 | 104 | def DenseNet161(**kwargs): 105 | return DenseNet(Bottleneck, [6,12,36,24], growth_rate=48, **kwargs) 106 | 107 | def densenet_cifar(*args, **kwargs): 108 | return DenseNet(Bottleneck, [6,12,24,16], growth_rate=12, **kwargs) 109 | 110 | densenet121 = DenseNet121 111 | densenet161 = DenseNet161 112 | densenet169 = DenseNet169 113 | densenet201 = DenseNet201 114 | 115 | def test(): 116 | net = densenet_cifar() 117 | x = torch.randn(1,3,32,32) 118 | y = net(x) 119 | print(y) 120 | -------------------------------------------------------------------------------- /robustness/cifar_models/inception.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | class BasicConv2d(nn.Module): 6 | 7 | def __init__(self, input_channels, output_channels, **kwargs): 8 | super().__init__() 9 | self.conv = nn.Conv2d(input_channels, output_channels, bias=False, **kwargs) 10 | self.bn = nn.BatchNorm2d(output_channels) 11 | self.relu = nn.ReLU(inplace=True) 12 | 13 | def forward(self, x): 14 | x = self.conv(x) 15 | x = self.bn(x) 16 | x = self.relu(x) 17 | 18 | return x 19 | 20 | #same naive inception module 21 | class InceptionA(nn.Module): 22 | 23 | def __init__(self, input_channels, pool_features): 24 | super().__init__() 25 | self.branch1x1 = BasicConv2d(input_channels, 64, kernel_size=1) 26 | 27 | self.branch5x5 = nn.Sequential( 28 | BasicConv2d(input_channels, 48, kernel_size=1), 29 | BasicConv2d(48, 64, kernel_size=5, padding=2) 30 | ) 31 | 32 | self.branch3x3 = nn.Sequential( 33 | BasicConv2d(input_channels, 64, kernel_size=1), 34 | BasicConv2d(64, 96, kernel_size=3, padding=1), 35 | BasicConv2d(96, 96, kernel_size=3, padding=1) 36 | ) 37 | 38 | self.branchpool = nn.Sequential( 39 | nn.AvgPool2d(kernel_size=3, stride=1, padding=1), 40 | BasicConv2d(input_channels, pool_features, kernel_size=3, padding=1) 41 | ) 42 | 43 | def forward(self, x): 44 | 45 | #x -> 1x1(same) 46 | branch1x1 = self.branch1x1(x) 47 | 48 | #x -> 1x1 -> 5x5(same) 49 | branch5x5 = self.branch5x5(x) 50 | #branch5x5 = self.branch5x5_2(branch5x5) 51 | 52 | #x -> 1x1 -> 3x3 -> 3x3(same) 53 | branch3x3 = self.branch3x3(x) 54 | 55 | #x -> pool -> 1x1(same) 56 | branchpool = self.branchpool(x) 57 | 58 | outputs = [branch1x1, branch5x5, branch3x3, branchpool] 59 | 60 | return torch.cat(outputs, 1) 61 | 62 | #downsample 63 | #Factorization into smaller convolutions 64 | class InceptionB(nn.Module): 65 | 66 | def __init__(self, input_channels): 67 | super().__init__() 68 | 69 | self.branch3x3 = BasicConv2d(input_channels, 384, kernel_size=3, stride=2) 70 | 71 | self.branch3x3stack = nn.Sequential( 72 | BasicConv2d(input_channels, 64, kernel_size=1), 73 | BasicConv2d(64, 96, kernel_size=3, padding=1), 74 | BasicConv2d(96, 96, kernel_size=3, stride=2) 75 | ) 76 | 77 | self.branchpool = nn.MaxPool2d(kernel_size=3, stride=2) 78 | 79 | def forward(self, x): 80 | 81 | #x - > 3x3(downsample) 82 | branch3x3 = self.branch3x3(x) 83 | 84 | #x -> 3x3 -> 3x3(downsample) 85 | branch3x3stack = self.branch3x3stack(x) 86 | 87 | #x -> avgpool(downsample) 88 | branchpool = self.branchpool(x) 89 | 90 | #"""We can use two parallel stride 2 blocks: P and C. P is a pooling 91 | #layer (either average or maximum pooling) the activation, both of 92 | #them are stride 2 the filter banks of which are concatenated as in 93 | #figure 10.""" 94 | outputs = [branch3x3, branch3x3stack, branchpool] 95 | 96 | return torch.cat(outputs, 1) 97 | 98 | #Factorizing Convolutions with Large Filter Size 99 | class InceptionC(nn.Module): 100 | def __init__(self, input_channels, channels_7x7): 101 | super().__init__() 102 | self.branch1x1 = BasicConv2d(input_channels, 192, kernel_size=1) 103 | 104 | c7 = channels_7x7 105 | 106 | #In theory, we could go even further and argue that one can replace any n x n 107 | #convolution by a 1 x n convolution followed by a n x 1 convolution and the 108 | #computational cost saving increases dramatically as n grows (see figure 6). 109 | self.branch7x7 = nn.Sequential( 110 | BasicConv2d(input_channels, c7, kernel_size=1), 111 | BasicConv2d(c7, c7, kernel_size=(7, 1), padding=(3, 0)), 112 | BasicConv2d(c7, 192, kernel_size=(1, 7), padding=(0, 3)) 113 | ) 114 | 115 | self.branch7x7stack = nn.Sequential( 116 | BasicConv2d(input_channels, c7, kernel_size=1), 117 | BasicConv2d(c7, c7, kernel_size=(7, 1), padding=(3, 0)), 118 | BasicConv2d(c7, c7, kernel_size=(1, 7), padding=(0, 3)), 119 | BasicConv2d(c7, c7, kernel_size=(7, 1), padding=(3, 0)), 120 | BasicConv2d(c7, 192, kernel_size=(1, 7), padding=(0, 3)) 121 | ) 122 | 123 | self.branch_pool = nn.Sequential( 124 | nn.AvgPool2d(kernel_size=3, stride=1, padding=1), 125 | BasicConv2d(input_channels, 192, kernel_size=1), 126 | ) 127 | 128 | def forward(self, x): 129 | 130 | #x -> 1x1(same) 131 | branch1x1 = self.branch1x1(x) 132 | 133 | #x -> 1layer 1*7 and 7*1 (same) 134 | branch7x7 = self.branch7x7(x) 135 | 136 | #x-> 2layer 1*7 and 7*1(same) 137 | branch7x7stack = self.branch7x7stack(x) 138 | 139 | #x-> avgpool (same) 140 | branchpool = self.branch_pool(x) 141 | 142 | outputs = [branch1x1, branch7x7, branch7x7stack, branchpool] 143 | 144 | return torch.cat(outputs, 1) 145 | 146 | class InceptionD(nn.Module): 147 | 148 | def __init__(self, input_channels): 149 | super().__init__() 150 | 151 | self.branch3x3 = nn.Sequential( 152 | BasicConv2d(input_channels, 192, kernel_size=1), 153 | BasicConv2d(192, 320, kernel_size=3, stride=2) 154 | ) 155 | 156 | self.branch7x7 = nn.Sequential( 157 | BasicConv2d(input_channels, 192, kernel_size=1), 158 | BasicConv2d(192, 192, kernel_size=(1, 7), padding=(0, 3)), 159 | BasicConv2d(192, 192, kernel_size=(7, 1), padding=(3, 0)), 160 | BasicConv2d(192, 192, kernel_size=3, stride=2) 161 | ) 162 | 163 | self.branchpool = nn.AvgPool2d(kernel_size=3, stride=2) 164 | 165 | def forward(self, x): 166 | 167 | #x -> 1x1 -> 3x3(downsample) 168 | branch3x3 = self.branch3x3(x) 169 | 170 | #x -> 1x1 -> 1x7 -> 7x1 -> 3x3 (downsample) 171 | branch7x7 = self.branch7x7(x) 172 | 173 | #x -> avgpool (downsample) 174 | branchpool = self.branchpool(x) 175 | 176 | outputs = [branch3x3, branch7x7, branchpool] 177 | 178 | return torch.cat(outputs, 1) 179 | 180 | 181 | #same 182 | class InceptionE(nn.Module): 183 | def __init__(self, input_channels): 184 | super().__init__() 185 | self.branch1x1 = BasicConv2d(input_channels, 320, kernel_size=1) 186 | 187 | self.branch3x3_1 = BasicConv2d(input_channels, 384, kernel_size=1) 188 | self.branch3x3_2a = BasicConv2d(384, 384, kernel_size=(1, 3), padding=(0, 1)) 189 | self.branch3x3_2b = BasicConv2d(384, 384, kernel_size=(3, 1), padding=(1, 0)) 190 | 191 | self.branch3x3stack_1 = BasicConv2d(input_channels, 448, kernel_size=1) 192 | self.branch3x3stack_2 = BasicConv2d(448, 384, kernel_size=3, padding=1) 193 | self.branch3x3stack_3a = BasicConv2d(384, 384, kernel_size=(1, 3), padding=(0, 1)) 194 | self.branch3x3stack_3b = BasicConv2d(384, 384, kernel_size=(3, 1), padding=(1, 0)) 195 | 196 | self.branch_pool = nn.Sequential( 197 | nn.AvgPool2d(kernel_size=3, stride=1, padding=1), 198 | BasicConv2d(input_channels, 192, kernel_size=1) 199 | ) 200 | 201 | def forward(self, x): 202 | 203 | #x -> 1x1 (same) 204 | branch1x1 = self.branch1x1(x) 205 | 206 | # x -> 1x1 -> 3x1 207 | # x -> 1x1 -> 1x3 208 | # concatenate(3x1, 1x3) 209 | #"""7. Inception modules with expanded the filter bank outputs. 210 | #This architecture is used on the coarsest (8 x 8) grids to promote 211 | #high dimensional representations, as suggested by principle 212 | #2 of Section 2.""" 213 | branch3x3 = self.branch3x3_1(x) 214 | branch3x3 = [ 215 | self.branch3x3_2a(branch3x3), 216 | self.branch3x3_2b(branch3x3) 217 | ] 218 | branch3x3 = torch.cat(branch3x3, 1) 219 | 220 | # x -> 1x1 -> 3x3 -> 1x3 221 | # x -> 1x1 -> 3x3 -> 3x1 222 | #concatenate(1x3, 3x1) 223 | branch3x3stack = self.branch3x3stack_1(x) 224 | branch3x3stack = self.branch3x3stack_2(branch3x3stack) 225 | branch3x3stack = [ 226 | self.branch3x3stack_3a(branch3x3stack), 227 | self.branch3x3stack_3b(branch3x3stack) 228 | ] 229 | branch3x3stack = torch.cat(branch3x3stack, 1) 230 | 231 | branchpool = self.branch_pool(x) 232 | 233 | outputs = [branch1x1, branch3x3, branch3x3stack, branchpool] 234 | 235 | return torch.cat(outputs, 1) 236 | 237 | class InceptionV3(nn.Module): 238 | 239 | def __init__(self, num_classes=10): 240 | super().__init__() 241 | self.Conv2d_1a_3x3 = BasicConv2d(3, 32, kernel_size=3, padding=1) 242 | self.Conv2d_2a_3x3 = BasicConv2d(32, 32, kernel_size=3, padding=1) 243 | self.Conv2d_2b_3x3 = BasicConv2d(32, 64, kernel_size=3, padding=1) 244 | self.Conv2d_3b_1x1 = BasicConv2d(64, 80, kernel_size=1) 245 | self.Conv2d_4a_3x3 = BasicConv2d(80, 192, kernel_size=3) 246 | 247 | #naive inception module 248 | self.Mixed_5b = InceptionA(192, pool_features=32) 249 | self.Mixed_5c = InceptionA(256, pool_features=64) 250 | self.Mixed_5d = InceptionA(288, pool_features=64) 251 | 252 | #downsample 253 | self.Mixed_6a = InceptionB(288) 254 | 255 | self.Mixed_6b = InceptionC(768, channels_7x7=128) 256 | self.Mixed_6c = InceptionC(768, channels_7x7=160) 257 | self.Mixed_6d = InceptionC(768, channels_7x7=160) 258 | self.Mixed_6e = InceptionC(768, channels_7x7=192) 259 | 260 | #downsample 261 | self.Mixed_7a = InceptionD(768) 262 | 263 | self.Mixed_7b = InceptionE(1280) 264 | self.Mixed_7c = InceptionE(2048) 265 | 266 | #6*6 feature size 267 | self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) 268 | self.dropout = nn.Dropout2d() 269 | self.linear = nn.Linear(2048, num_classes) 270 | 271 | def forward(self, x, with_latent=False, fake_relu=False, no_relu=False): 272 | assert (not fake_relu) and (not no_relu), \ 273 | "fake_relu and no_relu not yet supported for this architecture" 274 | #32 -> 30 275 | x = self.Conv2d_1a_3x3(x) 276 | x = self.Conv2d_2a_3x3(x) 277 | x = self.Conv2d_2b_3x3(x) 278 | x = self.Conv2d_3b_1x1(x) 279 | x = self.Conv2d_4a_3x3(x) 280 | 281 | #30 -> 30 282 | x = self.Mixed_5b(x) 283 | x = self.Mixed_5c(x) 284 | x = self.Mixed_5d(x) 285 | 286 | #30 -> 14 287 | #Efficient Grid Size Reduction to avoid representation 288 | #bottleneck 289 | x = self.Mixed_6a(x) 290 | 291 | #14 -> 14 292 | #"""In practice, we have found that employing this factorization does not 293 | #work well on early layers, but it gives very good results on medium 294 | #grid-sizes (On m x m feature maps, where m ranges between 12 and 20). 295 | #On that level, very good results can be achieved by using 1 x 7 convolutions 296 | #followed by 7 x 1 convolutions.""" 297 | x = self.Mixed_6b(x) 298 | x = self.Mixed_6c(x) 299 | x = self.Mixed_6d(x) 300 | x = self.Mixed_6e(x) 301 | 302 | #14 -> 6 303 | #Efficient Grid Size Reduction 304 | x = self.Mixed_7a(x) 305 | 306 | #6 -> 6 307 | #We are using this solution only on the coarsest grid, 308 | #since that is the place where producing high dimensional 309 | #sparse representation is the most critical as the ratio of 310 | #local processing (by 1 x 1 convolutions) is increased compared 311 | #to the spatial aggregation.""" 312 | x = self.Mixed_7b(x) 313 | x = self.Mixed_7c(x) 314 | 315 | #6 -> 1 316 | x = self.avgpool(x) 317 | x = self.dropout(x) 318 | latent = x.view(x.size(0), -1) 319 | out = self.linear(latent) 320 | if with_latent: 321 | return out, latent 322 | return out 323 | 324 | 325 | def inceptionv3(*args, **kwargs): 326 | return InceptionV3() 327 | -------------------------------------------------------------------------------- /robustness/cifar_models/resnet.py: -------------------------------------------------------------------------------- 1 | '''ResNet in PyTorch. 2 | For Pre-activation ResNet, see 'preact_resnet.py'. 3 | Reference: 4 | [1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun 5 | Deep Residual Learning for Image Recognition. arXiv:1512.03385 6 | ''' 7 | import torch 8 | import torch.nn as nn 9 | import torch.nn.functional as F 10 | from ..tools.custom_modules import SequentialWithArgs, FakeReLU 11 | 12 | class BasicBlock(nn.Module): 13 | expansion = 1 14 | 15 | def __init__(self, in_planes, planes, stride=1): 16 | super(BasicBlock, self).__init__() 17 | self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, 18 | padding=1, bias=False) 19 | self.bn1 = nn.BatchNorm2d(planes) 20 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, 21 | padding=1, bias=False) 22 | self.bn2 = nn.BatchNorm2d(planes) 23 | 24 | self.shortcut = nn.Sequential() 25 | if stride != 1 or in_planes != self.expansion*planes: 26 | self.shortcut = nn.Sequential( 27 | nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, 28 | stride=stride, bias=False), 29 | nn.BatchNorm2d(self.expansion*planes)) 30 | 31 | def forward(self, x, fake_relu=False): 32 | out = F.relu(self.bn1(self.conv1(x))) 33 | out = self.bn2(self.conv2(out)) 34 | out += self.shortcut(x) 35 | if fake_relu: 36 | return FakeReLU.apply(out) 37 | return F.relu(out) 38 | 39 | 40 | class Bottleneck(nn.Module): 41 | expansion = 4 42 | 43 | def __init__(self, in_planes, planes, stride=1): 44 | super(Bottleneck, self).__init__() 45 | self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False) 46 | self.bn1 = nn.BatchNorm2d(planes) 47 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, 48 | padding=1, bias=False) 49 | self.bn2 = nn.BatchNorm2d(planes) 50 | self.conv3 = nn.Conv2d(planes, self.expansion*planes, kernel_size=1, bias=False) 51 | self.bn3 = nn.BatchNorm2d(self.expansion*planes) 52 | 53 | self.shortcut = nn.Sequential() 54 | if stride != 1 or in_planes != self.expansion*planes: 55 | self.shortcut = nn.Sequential( 56 | nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False), 57 | nn.BatchNorm2d(self.expansion*planes) 58 | ) 59 | 60 | def forward(self, x, fake_relu=False): 61 | out = F.relu(self.bn1(self.conv1(x))) 62 | out = F.relu(self.bn2(self.conv2(out))) 63 | out = self.bn3(self.conv3(out)) 64 | out += self.shortcut(x) 65 | if fake_relu: 66 | return FakeReLU.apply(out) 67 | return F.relu(out) 68 | 69 | 70 | class ResNet(nn.Module): 71 | # feat_scale lets us deal with CelebA, other non-32x32 datasets 72 | def __init__(self, block, num_blocks, num_classes=10, feat_scale=1, wm=1): 73 | super(ResNet, self).__init__() 74 | 75 | widths = [64, 128, 256, 512] 76 | widths = [int(w * wm) for w in widths] 77 | 78 | self.in_planes = widths[0] 79 | self.conv1 = nn.Conv2d(3, self.in_planes, kernel_size=3, stride=1, 80 | padding=1, bias=False) 81 | self.bn1 = nn.BatchNorm2d(self.in_planes) 82 | self.layer1 = self._make_layer(block, widths[0], num_blocks[0], stride=1) 83 | self.layer2 = self._make_layer(block, widths[1], num_blocks[1], stride=2) 84 | self.layer3 = self._make_layer(block, widths[2], num_blocks[2], stride=2) 85 | self.layer4 = self._make_layer(block, widths[3], num_blocks[3], stride=2) 86 | self.linear = nn.Linear(feat_scale*widths[3]*block.expansion, num_classes) 87 | 88 | def _make_layer(self, block, planes, num_blocks, stride): 89 | strides = [stride] + [1]*(num_blocks-1) 90 | layers = [] 91 | for stride in strides: 92 | layers.append(block(self.in_planes, planes, stride)) 93 | self.in_planes = planes * block.expansion 94 | return SequentialWithArgs(*layers) 95 | 96 | def forward(self, x, with_latent=False, fake_relu=False, no_relu=False): 97 | assert (not no_relu), \ 98 | "no_relu not yet supported for this architecture" 99 | out = F.relu(self.bn1(self.conv1(x))) 100 | out = self.layer1(out) 101 | out = self.layer2(out) 102 | out = self.layer3(out) 103 | out = self.layer4(out, fake_relu=fake_relu) 104 | out = F.avg_pool2d(out, 4) 105 | pre_out = out.view(out.size(0), -1) 106 | final = self.linear(pre_out) 107 | if with_latent: 108 | return final, pre_out 109 | return final 110 | 111 | def ResNet18(**kwargs): 112 | return ResNet(BasicBlock, [2,2,2,2], **kwargs) 113 | 114 | def ResNet18Wide(**kwargs): 115 | return ResNet(BasicBlock, [2,2,2,2], wm=5, **kwargs) 116 | 117 | def ResNet18Thin(**kwargs): 118 | return ResNet(BasicBlock, [2,2,2,2], wd=.75, **kwargs) 119 | 120 | def ResNet34(**kwargs): 121 | return ResNet(BasicBlock, [3,4,6,3], **kwargs) 122 | 123 | def ResNet50(**kwargs): 124 | return ResNet(Bottleneck, [3,4,6,3], **kwargs) 125 | 126 | def ResNet101(**kwargs): 127 | return ResNet(Bottleneck, [3,4,23,3], **kwargs) 128 | 129 | def ResNet152(**kwargs): 130 | return ResNet(Bottleneck, [3,8,36,3], **kwargs) 131 | 132 | resnet50 = ResNet50 133 | resnet18 = ResNet18 134 | resnet34 = ResNet34 135 | resnet101 = ResNet101 136 | resnet152 = ResNet152 137 | resnet18wide = ResNet18Wide 138 | 139 | # resnet18thin = ResNet18Thin 140 | def test(): 141 | net = ResNet18() 142 | y = net(torch.randn(1,3,32,32)) 143 | print(y.size()) 144 | 145 | -------------------------------------------------------------------------------- /robustness/cifar_models/vgg.py: -------------------------------------------------------------------------------- 1 | '''VGG11/13/16/19 in Pytorch.''' 2 | import torch 3 | import torch.nn as nn 4 | 5 | cfg = { 6 | 'VGG11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 7 | 'VGG13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 8 | 'VGG16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'], 9 | 'VGG19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'], 10 | } 11 | 12 | class VGG(nn.Module): 13 | def __init__(self, vgg_name, num_classes=10): 14 | super(VGG, self).__init__() 15 | self.features = self._make_layers(cfg[vgg_name]) 16 | self.classifier = nn.Linear(512, num_classes) 17 | 18 | def forward(self, x, with_latent=False, fake_relu=False, no_relu=False): 19 | assert (not fake_relu) and (not no_relu), \ 20 | "fake_relu and no_relu not yet supported for this architecture" 21 | out = self.features(x) 22 | latent = out.view(out.size(0), -1) 23 | out = self.classifier(latent) 24 | if with_latent: 25 | return out, latent 26 | return out 27 | 28 | def _make_layers(self, cfg): 29 | layers = [] 30 | in_channels = 3 31 | for x in cfg: 32 | if x == 'M': 33 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] 34 | else: 35 | layers += [nn.Conv2d(in_channels, x, kernel_size=3, padding=1), 36 | nn.BatchNorm2d(x), 37 | nn.ReLU(inplace=True)] 38 | in_channels = x 39 | layers += [nn.AvgPool2d(kernel_size=1, stride=1)] 40 | return nn.Sequential(*layers) 41 | 42 | def VGG11(**kwargs): 43 | return VGG('VGG11', **kwargs) 44 | 45 | def VGG13(**kwargs): 46 | return VGG('VGG13', **kwargs) 47 | 48 | def VGG16(**kwargs): 49 | return VGG('VGG16', **kwargs) 50 | 51 | def VGG19(**kwargs): 52 | return VGG('VGG19', **kwargs) 53 | 54 | vgg11 = VGG11 55 | vgg13 = VGG13 56 | vgg16 = VGG16 57 | vgg19 = VGG19 58 | -------------------------------------------------------------------------------- /robustness/data_augmentation.py: -------------------------------------------------------------------------------- 1 | """ 2 | Module responsible for data augmentation constants and configuration. 3 | """ 4 | 5 | import torch as ch 6 | from torchvision import transforms 7 | 8 | # lighting transform 9 | # https://git.io/fhBOc 10 | IMAGENET_PCA = { 11 | 'eigval':ch.Tensor([0.2175, 0.0188, 0.0045]), 12 | 'eigvec':ch.Tensor([ 13 | [-0.5675, 0.7192, 0.4009], 14 | [-0.5808, -0.0045, -0.8140], 15 | [-0.5836, -0.6948, 0.4203], 16 | ]) 17 | } 18 | class Lighting(object): 19 | """ 20 | Lighting noise (see https://git.io/fhBOc) 21 | """ 22 | def __init__(self, alphastd, eigval, eigvec): 23 | self.alphastd = alphastd 24 | self.eigval = eigval 25 | self.eigvec = eigvec 26 | 27 | def __call__(self, img): 28 | if self.alphastd == 0: 29 | return img 30 | 31 | alpha = img.new().resize_(3).normal_(0, self.alphastd) 32 | rgb = self.eigvec.type_as(img).clone()\ 33 | .mul(alpha.view(1, 3).expand(3, 3))\ 34 | .mul(self.eigval.view(1, 3).expand(3, 3))\ 35 | .sum(1).squeeze() 36 | 37 | return img.add(rgb.view(3, 1, 1).expand_as(img)) 38 | 39 | # Special transforms for ImageNet(s) 40 | TRAIN_TRANSFORMS_IMAGENET = transforms.Compose([ 41 | transforms.RandomResizedCrop(224), 42 | transforms.RandomHorizontalFlip(), 43 | transforms.ColorJitter( 44 | brightness=0.1, 45 | contrast=0.1, 46 | saturation=0.1 47 | ), 48 | transforms.ToTensor(), 49 | Lighting(0.05, IMAGENET_PCA['eigval'], 50 | IMAGENET_PCA['eigvec']) 51 | ]) 52 | """ 53 | Standard training data augmentation for ImageNet-scale datasets: Random crop, 54 | Random flip, Color Jitter, and Lighting Transform (see https://git.io/fhBOc) 55 | """ 56 | 57 | TEST_TRANSFORMS_IMAGENET = transforms.Compose([ 58 | transforms.Resize(256), 59 | transforms.CenterCrop(224), 60 | transforms.ToTensor(), 61 | ]) 62 | """ 63 | Standard test data processing (no augmentation) for ImageNet-scale datasets, 64 | Resized to 256x256 then center cropped to 224x224. 65 | """ 66 | 67 | # Data Augmentation defaults 68 | TRAIN_TRANSFORMS_DEFAULT = lambda size: transforms.Compose([ 69 | transforms.RandomCrop(size, padding=4), 70 | transforms.RandomHorizontalFlip(), 71 | transforms.ColorJitter(.25,.25,.25), 72 | transforms.RandomRotation(2), 73 | transforms.ToTensor(), 74 | ]) 75 | """ 76 | Generic training data transform, given image side length does random cropping, 77 | flipping, color jitter, and rotation. Called as, for example, 78 | :meth:`robustness.data_augmentation.TRAIN_TRANSFORMS_DEFAULT(32)` for CIFAR-10. 79 | """ 80 | 81 | TEST_TRANSFORMS_DEFAULT = lambda size:transforms.Compose([ 82 | transforms.Resize(size), 83 | transforms.CenterCrop(size), 84 | transforms.ToTensor() 85 | ]) 86 | """ 87 | Generic test data transform (no augmentation) to complement 88 | :meth:`robustness.data_augmentation.TEST_TRANSFORMS_DEFAULT`, takes in an image 89 | side length. 90 | """ 91 | -------------------------------------------------------------------------------- /robustness/defaults.py: -------------------------------------------------------------------------------- 1 | """ 2 | This module is used to set up arguments and defaults. For information on how to 3 | use it, see Step 2 of the :doc:`../example_usage/training_lib_part_1` 4 | walkthrough. 5 | """ 6 | 7 | from . import attacker, datasets 8 | from .tools import helpers 9 | 10 | BY_DATASET = 'varies by dataset' 11 | REQ = 'REQUIRED' 12 | 13 | TRAINING_DEFAULTS = { 14 | datasets.CIFAR: { 15 | "epochs": 150, 16 | "batch_size": 128, 17 | "weight_decay":5e-4, 18 | "step_lr": 50 19 | }, 20 | datasets.CINIC: { 21 | "epochs": 150, 22 | "batch_size": 128, 23 | "weight_decay":5e-4, 24 | "step_lr": 50 25 | }, 26 | datasets.ImageNet: { 27 | "epochs": 200, 28 | "batch_size":256, 29 | "weight_decay":1e-4, 30 | "step_lr": 50 31 | }, 32 | datasets.Places365: { 33 | "epochs": 200, 34 | "batch_size":256, 35 | "weight_decay":1e-4, 36 | "step_lr": 50 37 | }, 38 | datasets.RestrictedImageNet: { 39 | "epochs": 150, 40 | "batch_size": 256, 41 | "weight_decay": 1e-4, 42 | "step_lr": 50 43 | }, 44 | datasets.CustomImageNet: { 45 | "epochs": 200, 46 | "batch_size": 256, 47 | "weight_decay": 1e-4, 48 | "step_lr": 50 49 | }, 50 | datasets.A2B: { 51 | "epochs": 150, 52 | "batch_size": 64, 53 | "weight_decay": 5e-4, 54 | "step_lr": 50 55 | }, 56 | datasets.OpenImages: { 57 | "epochs": 200, 58 | "batch_size":256, 59 | "weight_decay":1e-4, 60 | "step_lr": 50 61 | }, 62 | } 63 | """ 64 | Default hyperparameters for training by dataset (tested for resnet50). 65 | Parameters can be accessed as `TRAINING_DEFAULTS[dataset_class][param_name]` 66 | """ 67 | 68 | TRAINING_ARGS = [ 69 | ['out-dir', str, 'where to save training logs and checkpoints', REQ], 70 | ['epochs', int, 'number of epochs to train for', BY_DATASET], 71 | ['lr', float, 'initial learning rate for training', 0.1], 72 | ['weight-decay', float, 'SGD weight decay parameter', BY_DATASET], 73 | ['momentum', float, 'SGD momentum parameter', 0.9], 74 | ['step-lr', int, 'number of steps between step-lr-gamma x LR drops', BY_DATASET], 75 | ['step-lr-gamma', float, 'multiplier by which LR drops in step scheduler', 0.1], 76 | ['custom-lr-multiplier', str, 'LR multiplier sched (format: [(epoch, LR),...])', None], 77 | ['lr-interpolation', ["linear", "step"], 'Drop LR as step function or linearly', "step"], 78 | ['adv-train', [0, 1], 'whether to train adversarially', REQ], 79 | ['adv-eval', [0, 1], 'whether to adversarially evaluate', None], 80 | ['log-iters', int, 'how frequently (in epochs) to log', 5], 81 | ['save-ckpt-iters', int, 'how frequently (epochs) to save \ 82 | (-1 for none, only saves best and last)', -1] 83 | ] 84 | """ 85 | Arguments essential for the `train_model` function. 86 | 87 | *Format*: `[NAME, TYPE/CHOICES, HELP STRING, DEFAULT (REQ=required, 88 | BY_DATASET=looked up in TRAINING_DEFAULTS at runtime)]` 89 | """ 90 | 91 | PGD_ARGS = [ 92 | ['attack-steps', int, 'number of steps for PGD attack', 7], 93 | ['constraint', list(attacker.STEPS.keys()), 'adv constraint', REQ], 94 | ['eps', str , 'adversarial perturbation budget', REQ], 95 | ['attack-lr', str, 'step size for PGD', REQ], 96 | ['use-best', [0, 1], 'if 1 (0) use best (final) PGD step as example', 1], 97 | ['random-restarts', int, 'number of random PGD restarts for eval', 0], 98 | ['random-start', [0, 1], 'start with random noise instead of pgd step', 0], 99 | ['custom-eps-multiplier', str, 'eps mult. sched (same format as LR)', None] 100 | ] 101 | """ 102 | Arguments essential for the :meth:`robustness.train.train_model` function if 103 | adversarially training or evaluating. 104 | 105 | *Format*: `[NAME, TYPE/CHOICES, HELP STRING, DEFAULT (REQ=required, 106 | BY_DATASET=looked up in TRAINING_DEFAULTS at runtime)]` 107 | """ 108 | 109 | MODEL_LOADER_ARGS = [ 110 | ['dataset', list(datasets.DATASETS.keys()), '', REQ], 111 | ['data', str, 'path to the dataset', '/tmp/'], 112 | ['arch', str, 'architecture (see {cifar,imagenet}_models/', REQ], 113 | ['batch-size', int, 'batch size for data loading', BY_DATASET], 114 | ['workers', int, '# data loading workers', 30], 115 | ['resume', str, 'path to checkpoint to resume from', None], 116 | ['resume-optimizer', [0, 1], 'whether to also resume optimizers', 0], 117 | ['data-aug', [0, 1], 'whether to use data augmentation', 1], 118 | ['mixed-precision', [0, 1], 'whether to use MP training (faster)', 0], 119 | ] 120 | """ 121 | Arguments essential for constructing the model and dataloaders that will be fed 122 | into :meth:`robustness.train.train_model` or :meth:`robustness.train.eval_model` 123 | 124 | *Format*: `[NAME, TYPE/CHOICES, HELP STRING, DEFAULT (REQ=required, 125 | BY_DATASET=looked up in TRAINING_DEFAULTS at runtime)]` 126 | """ 127 | 128 | CONFIG_ARGS = [ 129 | ['config-path', str, 'config path for loading in parameters', None], 130 | ['eval-only', [0, 1], 'just run evaluation (no training)', 0], 131 | ['exp-name', str, 'where to save in (inside out_dir)', None] 132 | ] 133 | """ 134 | Arguments for main.py specifically 135 | 136 | *Format*: `[NAME, TYPE/CHOICES, HELP STRING, DEFAULT (REQ=required, 137 | BY_DATASET=looked up in TRAINING_DEFAULTS at runtime)]` 138 | """ 139 | 140 | def add_args_to_parser(arg_list, parser): 141 | """ 142 | Adds arguments from one of the argument lists above to a passed-in 143 | arparse.ArgumentParser object. Formats helpstrings according to the 144 | defaults, but does NOT set the actual argparse defaults (*important*). 145 | 146 | Args: 147 | arg_list (list) : A list of the same format as the lists above, i.e. 148 | containing entries of the form [NAME, TYPE/CHOICES, HELP, DEFAULT] 149 | parser (argparse.ArgumentParser) : An ArgumentParser object to which the 150 | arguments will be added 151 | 152 | Returns: 153 | The original parser, now with the arguments added in. 154 | """ 155 | for arg_name, arg_type, arg_help, arg_default in arg_list: 156 | has_choices = (type(arg_type) == list) 157 | kwargs = { 158 | 'type': type(arg_type[0]) if has_choices else arg_type, 159 | 'help': f"{arg_help} (default: {arg_default})" 160 | } 161 | if has_choices: kwargs['choices'] = arg_type 162 | parser.add_argument(f'--{arg_name}', **kwargs) 163 | return parser 164 | 165 | def check_and_fill_args(args, arg_list, ds_class): 166 | """ 167 | Fills in defaults based on an arguments list (e.g., TRAINING_ARGS) and a 168 | dataset class (e.g., datasets.CIFAR). 169 | 170 | Args: 171 | args (object) : Any object subclass exposing :samp:`setattr` and 172 | :samp:`getattr` (e.g. cox.utils.Parameters) 173 | arg_list (list) : A list of the same format as the lists above, i.e. 174 | containing entries of the form [NAME, TYPE/CHOICES, HELP, DEFAULT] 175 | ds_class (type) : A dataset class name (i.e. a 176 | :class:`robustness.datasets.DataSet` subclass name) 177 | 178 | Returns: 179 | args (object): The :samp:`args` object with all the defaults filled in according to :samp:`arg_list` defaults. 180 | """ 181 | for arg_name, _, _, arg_default in arg_list: 182 | name = arg_name.replace("-", "_") 183 | if helpers.has_attr(args, name): continue 184 | if arg_default == REQ: raise ValueError(f"{arg_name} required") 185 | elif arg_default == BY_DATASET: 186 | setattr(args, name, TRAINING_DEFAULTS[ds_class][name]) 187 | elif arg_default is not None: 188 | setattr(args, name, arg_default) 189 | return args 190 | 191 | 192 | 193 | -------------------------------------------------------------------------------- /robustness/imagenet_models/__init__.py: -------------------------------------------------------------------------------- 1 | from .resnet import * 2 | from .densenet import * 3 | from .vgg import * 4 | from .leaky_resnet import * 5 | from .alexnet import * 6 | from .squeezenet import * 7 | -------------------------------------------------------------------------------- /robustness/imagenet_models/alexnet.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from torch.hub import load_state_dict_from_url 3 | from ..tools.custom_modules import FakeReLUM 4 | 5 | __all__ = ['AlexNet', 'alexnet'] 6 | 7 | model_urls = { 8 | 'alexnet': 'https://download.pytorch.org/models/alexnet-owt-4df8aa71.pth', 9 | } 10 | 11 | class AlexNet(nn.Module): 12 | def __init__(self, num_classes=1000): 13 | super(AlexNet, self).__init__() 14 | self.features = nn.Sequential( 15 | nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2), 16 | nn.ReLU(inplace=True), 17 | nn.MaxPool2d(kernel_size=3, stride=2), 18 | nn.Conv2d(64, 192, kernel_size=5, padding=2), 19 | nn.ReLU(inplace=True), 20 | nn.MaxPool2d(kernel_size=3, stride=2), 21 | nn.Conv2d(192, 384, kernel_size=3, padding=1), 22 | nn.ReLU(inplace=True), 23 | nn.Conv2d(384, 256, kernel_size=3, padding=1), 24 | nn.ReLU(inplace=True), 25 | nn.Conv2d(256, 256, kernel_size=3, padding=1), 26 | nn.ReLU(inplace=True), 27 | nn.MaxPool2d(kernel_size=3, stride=2), 28 | ) 29 | self.avgpool = nn.AdaptiveAvgPool2d((6, 6)) 30 | self.classifier = nn.Sequential( 31 | nn.Dropout(), 32 | nn.Linear(256 * 6 * 6, 4096), 33 | nn.ReLU(inplace=True), 34 | nn.Dropout(), 35 | nn.Linear(4096, 4096), 36 | ) 37 | self.last_relu = nn.ReLU(inplace=True) 38 | self.last_relu_fake = FakeReLUM() 39 | self.last_layer = nn.Linear(4096, num_classes) 40 | 41 | def forward(self, x, with_latent=False, fake_relu=False, no_relu=False): 42 | x = self.features(x) 43 | x = self.avgpool(x) 44 | x = x.view(x.size(0), 256 * 6 * 6) 45 | x_latent = self.classifier(x) 46 | x_relu = self.last_relu_fake(x_latent) if fake_relu \ 47 | else self.last_relu(x_latent) 48 | x_out = self.last_layer(x_relu) 49 | 50 | if with_latent and no_relu: 51 | return x_out, x_latent 52 | if with_latent: 53 | return x_out, x_relu 54 | return x_out 55 | 56 | def alexnet(pretrained=False, progress=True, **kwargs): 57 | r"""AlexNet model architecture from the 58 | `"One weird trick..." `_ paper. 59 | 60 | Args: 61 | pretrained (bool): If True, returns a model pre-trained on ImageNet 62 | progress (bool): If True, displays a progress bar of the download to stderr 63 | """ 64 | model = AlexNet(**kwargs) 65 | if pretrained: 66 | state_dict = load_state_dict_from_url(model_urls['alexnet'], 67 | progress=progress) 68 | model.load_state_dict(state_dict) 69 | return model 70 | -------------------------------------------------------------------------------- /robustness/imagenet_models/densenet.py: -------------------------------------------------------------------------------- 1 | import re 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | import torch.utils.checkpoint as cp 6 | from collections import OrderedDict 7 | from torch.hub import load_state_dict_from_url 8 | from torch import Tensor 9 | from torch.jit.annotations import List 10 | from ..tools.custom_modules import FakeReLU 11 | 12 | 13 | __all__ = ['DenseNet', 'densenet121', 'densenet169', 'densenet201', 'densenet161'] 14 | 15 | model_urls = { 16 | 'densenet121': 'https://download.pytorch.org/models/densenet121-a639ec97.pth', 17 | 'densenet169': 'https://download.pytorch.org/models/densenet169-b2777c0a.pth', 18 | 'densenet201': 'https://download.pytorch.org/models/densenet201-c1103571.pth', 19 | 'densenet161': 'https://download.pytorch.org/models/densenet161-8d451a50.pth', 20 | } 21 | 22 | 23 | class _DenseLayer(nn.Module): 24 | def __init__(self, num_input_features, growth_rate, bn_size, drop_rate, memory_efficient=False): 25 | super(_DenseLayer, self).__init__() 26 | self.add_module('norm1', nn.BatchNorm2d(num_input_features)), 27 | self.add_module('relu1', nn.ReLU(inplace=True)), 28 | self.add_module('conv1', nn.Conv2d(num_input_features, bn_size * 29 | growth_rate, kernel_size=1, stride=1, 30 | bias=False)), 31 | self.add_module('norm2', nn.BatchNorm2d(bn_size * growth_rate)), 32 | self.add_module('relu2', nn.ReLU(inplace=True)), 33 | self.add_module('conv2', nn.Conv2d(bn_size * growth_rate, growth_rate, 34 | kernel_size=3, stride=1, padding=1, 35 | bias=False)), 36 | self.drop_rate = float(drop_rate) 37 | self.memory_efficient = memory_efficient 38 | 39 | def bn_function(self, inputs): 40 | # type: (List[Tensor]) -> Tensor 41 | concated_features = torch.cat(inputs, 1) 42 | bottleneck_output = self.conv1(self.relu1(self.norm1(concated_features))) # noqa: T484 43 | return bottleneck_output 44 | 45 | # todo: rewrite when torchscript supports any 46 | def any_requires_grad(self, input): 47 | # type: (List[Tensor]) -> bool 48 | for tensor in input: 49 | if tensor.requires_grad: 50 | return True 51 | return False 52 | 53 | # @torch.jit.unused # noqa: T484 54 | def call_checkpoint_bottleneck(self, input): 55 | # type: (List[Tensor]) -> Tensor 56 | def closure(*inputs): 57 | return self.bn_function(*inputs) 58 | 59 | return cp.checkpoint(closure, input) 60 | 61 | # @torch.jit._overload_method # noqa: F811 62 | def forward(self, input): 63 | # type: (List[Tensor]) -> (Tensor) 64 | pass 65 | 66 | # @torch.jit._overload_method # noqa: F811 67 | def forward(self, input): 68 | # type: (Tensor) -> (Tensor) 69 | pass 70 | 71 | # torchscript does not yet support *args, so we overload method 72 | # allowing it to take either a List[Tensor] or single Tensor 73 | def forward(self, input): # noqa: F811 74 | if isinstance(input, Tensor): 75 | prev_features = [input] 76 | else: 77 | prev_features = input 78 | 79 | if self.memory_efficient and self.any_requires_grad(prev_features): 80 | if torch.jit.is_scripting(): 81 | raise Exception("Memory Efficient not supported in JIT") 82 | 83 | bottleneck_output = self.call_checkpoint_bottleneck(prev_features) 84 | else: 85 | bottleneck_output = self.bn_function(prev_features) 86 | 87 | new_features = self.conv2(self.relu2(self.norm2(bottleneck_output))) 88 | if self.drop_rate > 0: 89 | new_features = F.dropout(new_features, p=self.drop_rate, 90 | training=self.training) 91 | return new_features 92 | 93 | 94 | class _DenseBlock(nn.ModuleDict): 95 | _version = 2 96 | 97 | def __init__(self, num_layers, num_input_features, bn_size, growth_rate, drop_rate, memory_efficient=False): 98 | super(_DenseBlock, self).__init__() 99 | for i in range(num_layers): 100 | layer = _DenseLayer( 101 | num_input_features + i * growth_rate, 102 | growth_rate=growth_rate, 103 | bn_size=bn_size, 104 | drop_rate=drop_rate, 105 | memory_efficient=memory_efficient, 106 | ) 107 | self.add_module('denselayer%d' % (i + 1), layer) 108 | 109 | def forward(self, init_features): 110 | features = [init_features] 111 | for name, layer in self.items(): 112 | new_features = layer(features) 113 | features.append(new_features) 114 | return torch.cat(features, 1) 115 | 116 | 117 | class _Transition(nn.Sequential): 118 | def __init__(self, num_input_features, num_output_features): 119 | super(_Transition, self).__init__() 120 | self.add_module('norm', nn.BatchNorm2d(num_input_features)) 121 | self.add_module('relu', nn.ReLU(inplace=True)) 122 | self.add_module('conv', nn.Conv2d(num_input_features, num_output_features, 123 | kernel_size=1, stride=1, bias=False)) 124 | self.add_module('pool', nn.AvgPool2d(kernel_size=2, stride=2)) 125 | 126 | 127 | class DenseNet(nn.Module): 128 | r"""Densenet-BC model class, based on 129 | `"Densely Connected Convolutional Networks" `_ 130 | 131 | Args: 132 | growth_rate (int) - how many filters to add each layer (`k` in paper) 133 | block_config (list of 4 ints) - how many layers in each pooling block 134 | num_init_features (int) - the number of filters to learn in the first convolution layer 135 | bn_size (int) - multiplicative factor for number of bottle neck layers 136 | (i.e. bn_size * k features in the bottleneck layer) 137 | drop_rate (float) - dropout rate after each dense layer 138 | num_classes (int) - number of classification classes 139 | memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient, 140 | but slower. Default: *False*. See `"paper" `_ 141 | """ 142 | 143 | __constants__ = ['features'] 144 | 145 | def __init__(self, growth_rate=32, block_config=(6, 12, 24, 16), 146 | num_init_features=64, bn_size=4, drop_rate=0, num_classes=1000, memory_efficient=False): 147 | 148 | super(DenseNet, self).__init__() 149 | 150 | # First convolution 151 | self.features = nn.Sequential(OrderedDict([ 152 | ('conv0', nn.Conv2d(3, num_init_features, kernel_size=7, stride=2, 153 | padding=3, bias=False)), 154 | ('norm0', nn.BatchNorm2d(num_init_features)), 155 | ('relu0', nn.ReLU(inplace=True)), 156 | ('pool0', nn.MaxPool2d(kernel_size=3, stride=2, padding=1)), 157 | ])) 158 | 159 | # Each denseblock 160 | num_features = num_init_features 161 | for i, num_layers in enumerate(block_config): 162 | block = _DenseBlock( 163 | num_layers=num_layers, 164 | num_input_features=num_features, 165 | bn_size=bn_size, 166 | growth_rate=growth_rate, 167 | drop_rate=drop_rate, 168 | memory_efficient=memory_efficient 169 | ) 170 | self.features.add_module('denseblock%d' % (i + 1), block) 171 | num_features = num_features + num_layers * growth_rate 172 | if i != len(block_config) - 1: 173 | trans = _Transition(num_input_features=num_features, 174 | num_output_features=num_features // 2) 175 | self.features.add_module('transition%d' % (i + 1), trans) 176 | num_features = num_features // 2 177 | 178 | # Final batch norm 179 | self.features.add_module('norm5', nn.BatchNorm2d(num_features)) 180 | 181 | # Linear layer 182 | self.classifier = nn.Linear(num_features, num_classes) 183 | 184 | # Official init from torch repo. 185 | for m in self.modules(): 186 | if isinstance(m, nn.Conv2d): 187 | nn.init.kaiming_normal_(m.weight) 188 | elif isinstance(m, nn.BatchNorm2d): 189 | nn.init.constant_(m.weight, 1) 190 | nn.init.constant_(m.bias, 0) 191 | elif isinstance(m, nn.Linear): 192 | nn.init.constant_(m.bias, 0) 193 | 194 | def forward(self, x, with_latent=False, fake_relu=False, no_relu=False): 195 | assert not no_relu, \ 196 | "DenseNet has no pre-ReLU activations, no_relu not supported" 197 | features = self.features(x) 198 | if fake_relu: 199 | out = FakeReLU.apply(features) 200 | else: 201 | out = F.relu(features, inplace=True) 202 | out = F.adaptive_avg_pool2d(out, (1, 1)) 203 | out = torch.flatten(out, 1) 204 | pre_out = out.clone() 205 | out = self.classifier(out) 206 | if with_latent: 207 | return out, pre_out 208 | else: 209 | return out 210 | 211 | def _load_state_dict(model, model_url, progress): 212 | # '.'s are no longer allowed in module names, but previous _DenseLayer 213 | # has keys 'norm.1', 'relu.1', 'conv.1', 'norm.2', 'relu.2', 'conv.2'. 214 | # They are also in the checkpoints in model_urls. This pattern is used 215 | # to find such keys. 216 | pattern = re.compile( 217 | r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$') 218 | 219 | state_dict = load_state_dict_from_url(model_url, progress=progress) 220 | for key in list(state_dict.keys()): 221 | res = pattern.match(key) 222 | if res: 223 | new_key = res.group(1) + res.group(2) 224 | state_dict[new_key] = state_dict[key] 225 | del state_dict[key] 226 | model.load_state_dict(state_dict) 227 | 228 | 229 | def _densenet(arch, growth_rate, block_config, num_init_features, pretrained, progress, 230 | **kwargs): 231 | model = DenseNet(growth_rate, block_config, num_init_features, **kwargs) 232 | if pretrained: 233 | _load_state_dict(model, model_urls[arch], progress) 234 | return model 235 | 236 | 237 | def densenet121(pretrained=False, progress=True, **kwargs): 238 | r"""Densenet-121 model from 239 | `"Densely Connected Convolutional Networks" `_ 240 | 241 | Args: 242 | pretrained (bool): If True, returns a model pre-trained on ImageNet 243 | progress (bool): If True, displays a progress bar of the download to stderr 244 | memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient, 245 | but slower. Default: *False*. See `"paper" `_ 246 | """ 247 | return _densenet('densenet121', 32, (6, 12, 24, 16), 64, pretrained, progress, 248 | **kwargs) 249 | 250 | 251 | 252 | def densenet161(pretrained=False, progress=True, **kwargs): 253 | r"""Densenet-161 model from 254 | `"Densely Connected Convolutional Networks" `_ 255 | 256 | Args: 257 | pretrained (bool): If True, returns a model pre-trained on ImageNet 258 | progress (bool): If True, displays a progress bar of the download to stderr 259 | memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient, 260 | but slower. Default: *False*. See `"paper" `_ 261 | """ 262 | return _densenet('densenet161', 48, (6, 12, 36, 24), 96, pretrained, progress, 263 | **kwargs) 264 | 265 | 266 | 267 | def densenet169(pretrained=False, progress=True, **kwargs): 268 | r"""Densenet-169 model from 269 | `"Densely Connected Convolutional Networks" `_ 270 | 271 | Args: 272 | pretrained (bool): If True, returns a model pre-trained on ImageNet 273 | progress (bool): If True, displays a progress bar of the download to stderr 274 | memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient, 275 | but slower. Default: *False*. See `"paper" `_ 276 | """ 277 | return _densenet('densenet169', 32, (6, 12, 32, 32), 64, pretrained, progress, 278 | **kwargs) 279 | 280 | 281 | 282 | def densenet201(pretrained=False, progress=True, **kwargs): 283 | r"""Densenet-201 model from 284 | `"Densely Connected Convolutional Networks" `_ 285 | 286 | Args: 287 | pretrained (bool): If True, returns a model pre-trained on ImageNet 288 | progress (bool): If True, displays a progress bar of the download to stderr 289 | memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient, 290 | but slower. Default: *False*. See `"paper" `_ 291 | """ 292 | return _densenet('densenet201', 32, (6, 12, 48, 32), 64, pretrained, progress, 293 | **kwargs) 294 | -------------------------------------------------------------------------------- /robustness/imagenet_models/leaky_resnet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.utils.model_zoo as model_zoo 4 | from ..tools.custom_modules import SequentialWithArgs, FakeReLU 5 | 6 | __all__ = ['leaky_resnet18', 'leaky_resnet34', 'leaky_resnet50', 7 | 'leaky_resnet101', 'leaky_resnet152'] 8 | 9 | def conv3x3(in_planes, out_planes, stride=1): 10 | """3x3 convolution with padding""" 11 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 12 | padding=1, bias=False) 13 | 14 | 15 | def conv1x1(in_planes, out_planes, stride=1): 16 | """1x1 convolution""" 17 | return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) 18 | 19 | 20 | class BasicBlock(nn.Module): 21 | expansion = 1 22 | 23 | def __init__(self, inplanes, planes, stride=1, downsample=None): 24 | super(BasicBlock, self).__init__() 25 | self.conv1 = conv3x3(inplanes, planes, stride) 26 | self.bn1 = nn.BatchNorm2d(planes) 27 | self.relu = nn.LeakyReLU(inplace=True) 28 | self.conv2 = conv3x3(planes, planes) 29 | self.bn2 = nn.BatchNorm2d(planes) 30 | self.downsample = downsample 31 | self.stride = stride 32 | 33 | def forward(self, x, fake_relu=False, no_relu=False): 34 | identity = x 35 | 36 | out = self.conv1(x) 37 | out = self.bn1(out) 38 | out = self.relu(out) 39 | 40 | out = self.conv2(out) 41 | out = self.bn2(out) 42 | 43 | if self.downsample is not None: 44 | identity = self.downsample(x) 45 | 46 | out += identity 47 | pre_out = out.clone() 48 | 49 | if fake_relu: 50 | return FakeReLU.apply(out) 51 | if no_relu: 52 | return out 53 | return self.relu(out) 54 | 55 | class Bottleneck(nn.Module): 56 | expansion = 4 57 | 58 | def __init__(self, inplanes, planes, stride=1, downsample=None): 59 | super(Bottleneck, self).__init__() 60 | self.conv1 = conv1x1(inplanes, planes) 61 | self.bn1 = nn.BatchNorm2d(planes) 62 | self.conv2 = conv3x3(planes, planes, stride) 63 | self.bn2 = nn.BatchNorm2d(planes) 64 | self.conv3 = conv1x1(planes, planes * self.expansion) 65 | self.bn3 = nn.BatchNorm2d(planes * self.expansion) 66 | self.relu = nn.LeakyReLU(inplace=True) 67 | self.downsample = downsample 68 | self.stride = stride 69 | 70 | def forward(self, x, fake_relu=False, no_relu=False): 71 | identity = x 72 | 73 | out = self.conv1(x) 74 | out = self.bn1(out) 75 | out = self.relu(out) 76 | 77 | out = self.conv2(out) 78 | out = self.bn2(out) 79 | out = self.relu(out) 80 | 81 | out = self.conv3(out) 82 | out = self.bn3(out) 83 | 84 | if self.downsample is not None: 85 | identity = self.downsample(x) 86 | 87 | out += identity 88 | 89 | if fake_relu: 90 | return FakeReLU.apply(out) 91 | if no_relu: 92 | return out 93 | 94 | return self.relu(out) 95 | 96 | class ResNet(nn.Module): 97 | def __init__(self, block, layers, num_classes=1000, zero_init_residual=False): 98 | super(ResNet, self).__init__() 99 | self.inplanes = 64 100 | self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, 101 | bias=False) 102 | self.bn1 = nn.BatchNorm2d(64) 103 | self.relu = nn.LeakyReLU(inplace=True) 104 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 105 | self.layer1 = self._make_layer(block, 64, layers[0]) 106 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2) 107 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2) 108 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2) 109 | self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) 110 | self.fc = nn.Linear(512 * block.expansion, num_classes) 111 | 112 | for m in self.modules(): 113 | if isinstance(m, nn.Conv2d): 114 | nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 115 | elif isinstance(m, nn.BatchNorm2d): 116 | nn.init.constant_(m.weight, 1) 117 | nn.init.constant_(m.bias, 0) 118 | 119 | # Zero-initialize the last BN in each residual branch, 120 | # so that the residual branch starts with zeros, and each residual block behaves like an identity. 121 | # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677 122 | if zero_init_residual: 123 | for m in self.modules(): 124 | if isinstance(m, Bottleneck): 125 | nn.init.constant_(m.bn3.weight, 0) 126 | elif isinstance(m, BasicBlock): 127 | nn.init.constant_(m.bn2.weight, 0) 128 | 129 | def _make_layer(self, block, planes, blocks, stride=1): 130 | downsample = None 131 | if stride != 1 or self.inplanes != planes * block.expansion: 132 | downsample = nn.Sequential( 133 | conv1x1(self.inplanes, planes * block.expansion, stride), 134 | nn.BatchNorm2d(planes * block.expansion), 135 | ) 136 | 137 | layers = [] 138 | layers.append(block(self.inplanes, planes, stride, downsample)) 139 | self.inplanes = planes * block.expansion 140 | for _ in range(1, blocks): 141 | layers.append(block(self.inplanes, planes)) 142 | 143 | return SequentialWithArgs(*layers) 144 | 145 | def forward(self, x, with_latent=False, fake_relu=False, no_relu=False): 146 | x = self.conv1(x) 147 | x = self.bn1(x) 148 | x = self.relu(x) 149 | x = self.maxpool(x) 150 | 151 | x = self.layer1(x) 152 | x = self.layer2(x) 153 | x = self.layer3(x) 154 | x = self.layer4(x, fake_relu=fake_relu, no_relu=no_relu) 155 | 156 | x = self.avgpool(x) 157 | pre_out = x.view(x.size(0), -1) 158 | final = self.fc(pre_out) 159 | if with_latent: 160 | return final, pre_out 161 | return final 162 | 163 | def leaky_resnet18(pretrained=False, **kwargs): 164 | """Constructs a ResNet-18 model. 165 | 166 | Args: 167 | pretrained (bool): If True, returns a model pre-trained on ImageNet 168 | """ 169 | model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs) 170 | if pretrained: raise NotImplementedError 171 | return model 172 | 173 | 174 | def leaky_resnet34(pretrained=False, **kwargs): 175 | """Constructs a ResNet-34 model. 176 | 177 | Args: 178 | pretrained (bool): If True, returns a model pre-trained on ImageNet 179 | """ 180 | model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs) 181 | if pretrained: raise NotImplementedError 182 | return model 183 | 184 | 185 | def leaky_resnet50(pretrained=False, **kwargs): 186 | """Constructs a ResNet-50 model. 187 | 188 | Args: 189 | pretrained (bool): If True, returns a model pre-trained on ImageNet 190 | """ 191 | model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs) 192 | if pretrained: raise NotImplementedError 193 | return model 194 | 195 | 196 | def leaky_resnet101(pretrained=False, **kwargs): 197 | """Constructs a ResNet-101 model. 198 | 199 | Args: 200 | pretrained (bool): If True, returns a model pre-trained on ImageNet 201 | """ 202 | model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs) 203 | if pretrained: raise NotImplementedError 204 | return model 205 | 206 | 207 | def leaky_resnet152(pretrained=False, **kwargs): 208 | """Constructs a ResNet-152 model. 209 | 210 | Args: 211 | pretrained (bool): If True, returns a model pre-trained on ImageNet 212 | """ 213 | model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs) 214 | if pretrained: raise NotImplementedError 215 | return model 216 | -------------------------------------------------------------------------------- /robustness/imagenet_models/squeezenet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.init as init 4 | from torch.hub import load_state_dict_from_url 5 | from ..tools.custom_modules import FakeReLUM 6 | 7 | 8 | __all__ = ['SqueezeNet', 'squeezenet1_0', 'squeezenet1_1'] 9 | 10 | model_urls = { 11 | 'squeezenet1_0': 'https://download.pytorch.org/models/squeezenet1_0-a815701f.pth', 12 | 'squeezenet1_1': 'https://download.pytorch.org/models/squeezenet1_1-f364aa15.pth', 13 | } 14 | 15 | 16 | class Fire(nn.Module): 17 | 18 | def __init__(self, inplanes, squeeze_planes, 19 | expand1x1_planes, expand3x3_planes): 20 | super(Fire, self).__init__() 21 | self.inplanes = inplanes 22 | self.squeeze = nn.Conv2d(inplanes, squeeze_planes, kernel_size=1) 23 | self.squeeze_activation = nn.ReLU(inplace=True) 24 | self.expand1x1 = nn.Conv2d(squeeze_planes, expand1x1_planes, 25 | kernel_size=1) 26 | self.expand1x1_activation = nn.ReLU(inplace=True) 27 | self.expand3x3 = nn.Conv2d(squeeze_planes, expand3x3_planes, 28 | kernel_size=3, padding=1) 29 | self.expand3x3_activation = nn.ReLU(inplace=True) 30 | 31 | def forward(self, x): 32 | x = self.squeeze_activation(self.squeeze(x)) 33 | return torch.cat([ 34 | self.expand1x1_activation(self.expand1x1(x)), 35 | self.expand3x3_activation(self.expand3x3(x)) 36 | ], 1) 37 | 38 | 39 | class SqueezeNet(nn.Module): 40 | 41 | def __init__(self, version='1_0', num_classes=1000): 42 | super(SqueezeNet, self).__init__() 43 | self.num_classes = num_classes 44 | if version == '1_0': 45 | self.features = nn.Sequential( 46 | nn.Conv2d(3, 96, kernel_size=7, stride=2), 47 | nn.ReLU(inplace=True), 48 | nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), 49 | Fire(96, 16, 64, 64), 50 | Fire(128, 16, 64, 64), 51 | Fire(128, 32, 128, 128), 52 | nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), 53 | Fire(256, 32, 128, 128), 54 | Fire(256, 48, 192, 192), 55 | Fire(384, 48, 192, 192), 56 | Fire(384, 64, 256, 256), 57 | nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), 58 | Fire(512, 64, 256, 256), 59 | ) 60 | elif version == '1_1': 61 | self.features = nn.Sequential( 62 | nn.Conv2d(3, 64, kernel_size=3, stride=2), 63 | nn.ReLU(inplace=True), 64 | nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), 65 | Fire(64, 16, 64, 64), 66 | Fire(128, 16, 64, 64), 67 | nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), 68 | Fire(128, 32, 128, 128), 69 | Fire(256, 32, 128, 128), 70 | nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), 71 | Fire(256, 48, 192, 192), 72 | Fire(384, 48, 192, 192), 73 | Fire(384, 64, 256, 256), 74 | Fire(512, 64, 256, 256), 75 | ) 76 | else: 77 | raise ValueError("Unsupported SqueezeNet version {version}:" 78 | "1_0 or 1_1 expected".format(version=version)) 79 | 80 | # Final convolution is initialized differently from the rest 81 | final_conv = nn.Conv2d(512, self.num_classes, kernel_size=1) 82 | self.classifier = nn.Sequential( 83 | nn.Dropout(p=0.5), 84 | final_conv, 85 | nn.ReLU(inplace=True), 86 | nn.AdaptiveAvgPool2d((1, 1)) 87 | ) 88 | self.last_relu = nn.ReLU(inplace=True) 89 | self.last_relu_fake = FakeReLUM() 90 | 91 | for m in self.modules(): 92 | if isinstance(m, nn.Conv2d): 93 | if m is final_conv: 94 | init.normal_(m.weight, mean=0.0, std=0.01) 95 | else: 96 | init.kaiming_uniform_(m.weight) 97 | if m.bias is not None: 98 | init.constant_(m.bias, 0) 99 | 100 | def forward(self, x, with_latent=False, fake_relu=False, no_relu=False): 101 | x = self.features(x) 102 | x_latent = self.classifier[:2](x) 103 | x_relu = self.last_relu(x_latent) if not fake_relu else self.classifier_last_relu_fake(x_latent) 104 | x_out = self.classifier[-1:](x_relu) 105 | x_out = torch.flatten(x_out, 1) 106 | 107 | if with_latent and no_relu: 108 | return x_out, x_latent # potentially will need to flatten x_latent 109 | if with_latent: 110 | return x_out, x_relu # potentially will need to flatten x_relu 111 | return x_out 112 | 113 | 114 | def _squeezenet(version, pretrained, progress, **kwargs): 115 | model = SqueezeNet(version, **kwargs) 116 | if pretrained: 117 | arch = 'squeezenet' + version 118 | state_dict = load_state_dict_from_url(model_urls[arch], 119 | progress=progress) 120 | model.load_state_dict(state_dict) 121 | return model 122 | 123 | 124 | def squeezenet1_0(pretrained=False, progress=True, **kwargs): 125 | r"""SqueezeNet model architecture from the `"SqueezeNet: AlexNet-level 126 | accuracy with 50x fewer parameters and <0.5MB model size" 127 | `_ paper. 128 | 129 | Args: 130 | pretrained (bool): If True, returns a model pre-trained on ImageNet 131 | progress (bool): If True, displays a progress bar of the download to stderr 132 | """ 133 | return _squeezenet('1_0', pretrained, progress, **kwargs) 134 | 135 | 136 | 137 | def squeezenet1_1(pretrained=False, progress=True, **kwargs): 138 | r"""SqueezeNet 1.1 model from the `official SqueezeNet repo 139 | `_. 140 | SqueezeNet 1.1 has 2.4x less computation and slightly fewer parameters 141 | than SqueezeNet 1.0, without sacrificing accuracy. 142 | 143 | Args: 144 | pretrained (bool): If True, returns a model pre-trained on ImageNet 145 | progress (bool): If True, displays a progress bar of the download to stderr 146 | """ 147 | return _squeezenet('1_1', pretrained, progress, **kwargs) 148 | -------------------------------------------------------------------------------- /robustness/imagenet_models/vgg.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch 3 | from torch.hub import load_state_dict_from_url 4 | from ..tools.custom_modules import FakeReLUM 5 | 6 | __all__ = [ 7 | 'VGG', 'vgg11', 'vgg11_bn', 'vgg13', 'vgg13_bn', 'vgg16', 'vgg16_bn', 8 | 'vgg19_bn', 'vgg19', 9 | ] 10 | 11 | 12 | model_urls = { 13 | 'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth', 14 | 'vgg13': 'https://download.pytorch.org/models/vgg13-c768596a.pth', 15 | 'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth', 16 | 'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth', 17 | 'vgg11_bn': 'https://download.pytorch.org/models/vgg11_bn-6002323d.pth', 18 | 'vgg13_bn': 'https://download.pytorch.org/models/vgg13_bn-abd245e5.pth', 19 | 'vgg16_bn': 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth', 20 | 'vgg19_bn': 'https://download.pytorch.org/models/vgg19_bn-c79401a0.pth', 21 | } 22 | 23 | class VGG(nn.Module): 24 | def __init__(self, features, num_classes=1000, init_weights=True): 25 | super(VGG, self).__init__() 26 | self.features = features 27 | self.avgpool = nn.AdaptiveAvgPool2d((7, 7)) 28 | self.classifier = nn.Sequential( 29 | nn.Linear(512 * 7 * 7, 4096), 30 | nn.ReLU(), 31 | nn.Dropout(), 32 | nn.Linear(4096, 4096), 33 | nn.ReLU(), 34 | nn.Dropout(), 35 | nn.Linear(4096, num_classes) 36 | ) 37 | self.last_relu = nn.ReLU() 38 | self.last_relu_fake = FakeReLUM() 39 | if init_weights: 40 | self._initialize_weights() 41 | 42 | def forward(self, x, with_latent=False, fake_relu=False, no_relu=False): 43 | feats = self.features(x) 44 | pooled = self.avgpool(feats) 45 | x = pooled.view(pooled.size(0), -1) 46 | x_latent = self.classifier[:4](x) 47 | x_relu = self.last_relu_fake(x_latent) if fake_relu \ 48 | else self.last_relu(x_latent) 49 | x_out = self.classifier[-2:](x_relu) 50 | 51 | if with_latent and no_relu: 52 | return x_out, x_latent 53 | if with_latent: 54 | return x_out, x_relu 55 | return x_out 56 | 57 | def _initialize_weights(self): 58 | for m in self.modules(): 59 | if isinstance(m, nn.Conv2d): 60 | nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 61 | if m.bias is not None: 62 | nn.init.constant_(m.bias, 0) 63 | elif isinstance(m, nn.BatchNorm2d): 64 | nn.init.constant_(m.weight, 1) 65 | nn.init.constant_(m.bias, 0) 66 | elif isinstance(m, nn.Linear): 67 | nn.init.normal_(m.weight, 0, 0.01) 68 | nn.init.constant_(m.bias, 0) 69 | 70 | def make_layers(cfg, batch_norm=False): 71 | layers = [] 72 | in_channels = 3 73 | for v in cfg: 74 | if v == 'M': 75 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] 76 | else: 77 | conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1) 78 | if batch_norm: 79 | layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU()] 80 | else: 81 | layers += [conv2d, nn.ReLU()] 82 | in_channels = v 83 | return nn.Sequential(*layers) 84 | 85 | 86 | cfgs = { 87 | 'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 88 | 'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 89 | 'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'], 90 | 'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'], 91 | } 92 | 93 | 94 | def _vgg(arch, cfg, batch_norm, pretrained, progress, **kwargs): 95 | if pretrained: 96 | kwargs['init_weights'] = False 97 | model = VGG(make_layers(cfgs[cfg], batch_norm=batch_norm), **kwargs) 98 | if pretrained: 99 | state_dict = load_state_dict_from_url(model_urls[arch], 100 | progress=progress) 101 | model.load_state_dict(state_dict) 102 | return model 103 | 104 | 105 | def vgg11(pretrained=False, progress=True, **kwargs): 106 | r"""VGG 11-layer model (configuration "A") from 107 | `"Very Deep Convolutional Networks For Large-Scale Image Recognition" '_ 108 | 109 | Args: 110 | pretrained (bool): If True, returns a model pre-trained on ImageNet 111 | progress (bool): If True, displays a progress bar of the download to stderr 112 | """ 113 | return _vgg('vgg11', 'A', False, pretrained, progress, **kwargs) 114 | 115 | 116 | def vgg11_bn(pretrained=False, progress=True, **kwargs): 117 | r"""VGG 11-layer model (configuration "A") with batch normalization 118 | `"Very Deep Convolutional Networks For Large-Scale Image Recognition" '_ 119 | 120 | Args: 121 | pretrained (bool): If True, returns a model pre-trained on ImageNet 122 | progress (bool): If True, displays a progress bar of the download to stderr 123 | """ 124 | return _vgg('vgg11_bn', 'A', True, pretrained, progress, **kwargs) 125 | 126 | 127 | def vgg13(pretrained=False, progress=True, **kwargs): 128 | r"""VGG 13-layer model (configuration "B") 129 | `"Very Deep Convolutional Networks For Large-Scale Image Recognition" '_ 130 | 131 | Args: 132 | pretrained (bool): If True, returns a model pre-trained on ImageNet 133 | progress (bool): If True, displays a progress bar of the download to stderr 134 | """ 135 | return _vgg('vgg13', 'B', False, pretrained, progress, **kwargs) 136 | 137 | 138 | def vgg13_bn(pretrained=False, progress=True, **kwargs): 139 | r"""VGG 13-layer model (configuration "B") with batch normalization 140 | `"Very Deep Convolutional Networks For Large-Scale Image Recognition" '_ 141 | 142 | Args: 143 | pretrained (bool): If True, returns a model pre-trained on ImageNet 144 | progress (bool): If True, displays a progress bar of the download to stderr 145 | """ 146 | return _vgg('vgg13_bn', 'B', True, pretrained, progress, **kwargs) 147 | 148 | 149 | def vgg16(pretrained=False, progress=True, **kwargs): 150 | r"""VGG 16-layer model (configuration "D") 151 | `"Very Deep Convolutional Networks For Large-Scale Image Recognition" '_ 152 | 153 | Args: 154 | pretrained (bool): If True, returns a model pre-trained on ImageNet 155 | progress (bool): If True, displays a progress bar of the download to stderr 156 | """ 157 | return _vgg('vgg16', 'D', False, pretrained, progress, **kwargs) 158 | 159 | 160 | def vgg16_bn(pretrained=False, progress=True, **kwargs): 161 | r"""VGG 16-layer model (configuration "D") with batch normalization 162 | `"Very Deep Convolutional Networks For Large-Scale Image Recognition" '_ 163 | 164 | Args: 165 | pretrained (bool): If True, returns a model pre-trained on ImageNet 166 | progress (bool): If True, displays a progress bar of the download to stderr 167 | """ 168 | return _vgg('vgg16_bn', 'D', True, pretrained, progress, **kwargs) 169 | 170 | 171 | def vgg19(pretrained=False, progress=True, **kwargs): 172 | r"""VGG 19-layer model (configuration "E") 173 | `"Very Deep Convolutional Networks For Large-Scale Image Recognition" '_ 174 | 175 | Args: 176 | pretrained (bool): If True, returns a model pre-trained on ImageNet 177 | progress (bool): If True, displays a progress bar of the download to stderr 178 | """ 179 | return _vgg('vgg19', 'E', False, pretrained, progress, **kwargs) 180 | 181 | 182 | def vgg19_bn(pretrained=False, progress=True, **kwargs): 183 | r"""VGG 19-layer model (configuration 'E') with batch normalization 184 | `"Very Deep Convolutional Networks For Large-Scale Image Recognition" '_ 185 | 186 | Args: 187 | pretrained (bool): If True, returns a model pre-trained on ImageNet 188 | progress (bool): If True, displays a progress bar of the download to stderr 189 | """ 190 | return _vgg('vgg19_bn', 'E', True, pretrained, progress, **kwargs) 191 | -------------------------------------------------------------------------------- /robustness/loaders.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | from . import cifar_models 4 | from .tools import folder 5 | 6 | import os 7 | if int(os.environ.get("NOTEBOOK_MODE", 0)) == 1: 8 | from tqdm import tqdm_notebook as tqdm 9 | else: 10 | from tqdm import tqdm as tqdm 11 | 12 | import shutil 13 | import time 14 | import numpy as np 15 | import torch as ch 16 | import torch.utils.data 17 | from torch.utils.data import DataLoader 18 | from torch.utils.data import Subset 19 | import torchvision.transforms as transforms 20 | from torch.utils.data import DataLoader 21 | from . import imagenet_models as models 22 | 23 | def make_loaders(workers, batch_size, transforms, data_path, data_aug=True, 24 | custom_class=None, dataset="", label_mapping=None, subset=None, 25 | subset_type='rand', subset_start=0, val_batch_size=None, 26 | only_val=False, shuffle_train=True, shuffle_val=True, seed=1, 27 | custom_class_args=None): 28 | ''' 29 | **INTERNAL FUNCTION** 30 | 31 | This is an internal function that makes a loader for any dataset. You 32 | probably want to call dataset.make_loaders for a specific dataset, 33 | which only requires workers and batch_size. For example: 34 | 35 | >>> cifar_dataset = CIFAR10('/path/to/cifar') 36 | >>> train_loader, val_loader = cifar_dataset.make_loaders(workers=10, batch_size=128) 37 | >>> # train_loader and val_loader are just PyTorch dataloaders 38 | ''' 39 | print(f"==> Preparing dataset {dataset}..") 40 | transform_train, transform_test = transforms 41 | if not data_aug: 42 | transform_train = transform_test 43 | 44 | if not val_batch_size: 45 | val_batch_size = batch_size 46 | 47 | if not custom_class: 48 | train_path = os.path.join(data_path, 'train') 49 | test_path = os.path.join(data_path, 'val') 50 | if not os.path.exists(test_path): 51 | test_path = os.path.join(data_path, 'test') 52 | 53 | if not os.path.exists(test_path): 54 | raise ValueError("Test data must be stored in dataset/test or {0}".format(test_path)) 55 | 56 | if not only_val: 57 | train_set = folder.ImageFolder(root=train_path, transform=transform_train, 58 | label_mapping=label_mapping) 59 | test_set = folder.ImageFolder(root=test_path, transform=transform_test, 60 | label_mapping=label_mapping) 61 | else: 62 | if custom_class_args is None: custom_class_args = {} 63 | if not only_val: 64 | train_set = custom_class(root=data_path, train=True, download=True, 65 | transform=transform_train, **custom_class_args) 66 | test_set = custom_class(root=data_path, train=False, download=True, 67 | transform=transform_test, **custom_class_args) 68 | 69 | if not only_val: 70 | attrs = ["samples", "train_data", "data"] 71 | vals = {attr: hasattr(train_set, attr) for attr in attrs} 72 | assert any(vals.values()), f"dataset must expose one of {attrs}" 73 | train_sample_count = len(getattr(train_set,[k for k in vals if vals[k]][0])) 74 | 75 | if (not only_val) and (subset is not None) and (subset <= train_sample_count): 76 | assert not only_val 77 | if subset_type == 'rand': 78 | rng = np.random.RandomState(seed) 79 | subset = rng.choice(list(range(train_sample_count)), size=subset+subset_start, replace=False) 80 | subset = subset[subset_start:] 81 | elif subset_type == 'first': 82 | subset = np.arange(subset_start, subset_start + subset) 83 | else: 84 | subset = np.arange(train_sample_count - subset, train_sample_count) 85 | 86 | train_set = Subset(train_set, subset) 87 | 88 | if not only_val: 89 | train_loader = DataLoader(train_set, batch_size=batch_size, 90 | shuffle=shuffle_train, num_workers=workers, pin_memory=True) 91 | 92 | test_loader = DataLoader(test_set, batch_size=val_batch_size, 93 | shuffle=shuffle_val, num_workers=workers, pin_memory=True) 94 | 95 | if only_val: 96 | return None, test_loader 97 | 98 | return train_loader, test_loader 99 | 100 | ## loader wrapper (for adding custom functions to dataloader) 101 | class PerEpochLoader: 102 | ''' 103 | A blend between TransformedLoader and LambdaLoader: stores the whole loader 104 | in memory, but recomputes it from scratch every epoch, instead of just once 105 | at initialization. 106 | ''' 107 | def __init__(self, loader, func, do_tqdm=True): 108 | self.orig_loader = loader 109 | self.func = func 110 | self.do_tqdm = do_tqdm 111 | self.data_loader = self.compute_loader() 112 | self.loader = iter(self.data_loader) 113 | 114 | def compute_loader(self): 115 | return TransformedLoader(self.orig_loader, self.func, None, 116 | self.orig_loader.num_workers, self.orig_loader.batch_size, 117 | do_tqdm=self.do_tqdm) 118 | 119 | def __len__(self): 120 | return len(self.orig_loader) 121 | 122 | def __getattr__(self, attr): 123 | return getattr(self.data_loader, attr) 124 | 125 | def __iter__(self): 126 | return self 127 | 128 | def __next__(self): 129 | try: 130 | return next(self.loader) 131 | except StopIteration as e: 132 | self.data_loader = self.compute_loader() 133 | self.loader = iter(self.data_loader) 134 | raise StopIteration 135 | 136 | return self.func(im, targ) 137 | 138 | class LambdaLoader: 139 | ''' 140 | This is a class that allows one to apply any given (fixed) 141 | transformation to the output from the loader in *real-time*. 142 | 143 | For instance, you could use for applications such as custom 144 | data augmentation and adding image/label noise. 145 | 146 | Note that the LambdaLoader is the final transformation that 147 | is applied to image-label pairs from the dataset as part of the 148 | loading process---i.e., other (standard) transformations such 149 | as data augmentation can only be applied *before* passing the 150 | data through the LambdaLoader. 151 | 152 | For more information see :ref:`our detailed walkthrough ` 153 | 154 | ''' 155 | 156 | def __init__(self, loader, func): 157 | ''' 158 | Args: 159 | loader (PyTorch dataloader) : loader for dataset (*required*). 160 | func (function) : fixed transformation to be applied to 161 | every batch in real-time (*required*). It takes in 162 | (images, labels) and returns (images, labels) of the 163 | same shape. 164 | ''' 165 | self.data_loader = loader 166 | self.loader = iter(self.data_loader) 167 | self.func = func 168 | 169 | def __len__(self): 170 | return len(self.data_loader) 171 | 172 | def __iter__(self): 173 | return self 174 | 175 | def __getattr__(self, attr): 176 | return getattr(self.data_loader, attr) 177 | 178 | def __next__(self): 179 | try: 180 | im, targ = next(self.loader) 181 | except StopIteration as e: 182 | self.loader = iter(self.data_loader) 183 | raise StopIteration 184 | 185 | return self.func(im, targ) 186 | 187 | def __getattr__(self, attr): 188 | return getattr(self.data_loader, attr) 189 | 190 | def TransformedLoader(loader, func, transforms, workers=None, 191 | batch_size=None, do_tqdm=False, augment=False, fraction=1.0, 192 | shuffle=True): 193 | ''' 194 | This is a function that allows one to apply any given (fixed) 195 | transformation to the output from the loader *once*. 196 | 197 | For instance, you could use for applications such as assigning 198 | random labels to all the images (before training). 199 | 200 | The TransformedLoader also supports the application of addiotional 201 | transformations (such as standard data augmentation) after the fixed 202 | function. 203 | 204 | For more information see :ref:`our detailed walkthrough ` 205 | 206 | Args: 207 | loader (PyTorch dataloader) : loader for dataset 208 | func (function) : fixed transformation to be applied once. It takes 209 | in (images, labels) and returns (images, labels) with the same shape 210 | in every dimension except for the first, i.e., batch dimension 211 | (which can be any length). 212 | transforms (torchvision.transforms) : transforms to apply 213 | to the training images from the dataset (after func) (*required*). 214 | workers (int) : number of workers for data fetching (*required*). 215 | batch_size (int) : batch size for the data loaders (*required*). 216 | do_tqdm (bool) : if True, show a tqdm progress bar for the attack. 217 | augment (bool) : if True, the output loader contains both the original 218 | (untransformed), and new transformed image-label pairs. 219 | fraction (float): fraction of image-label pairs in the output loader 220 | which are transformed. The remainder is just original image-label 221 | pairs from loader. 222 | shuffle (bool) : whether or not the resulting loader should shuffle every 223 | epoch (defaults to True) 224 | 225 | Returns: 226 | A loader and validation loader according to the 227 | parameters given. These are standard PyTorch data loaders, and 228 | thus can just be used via: 229 | 230 | >>> output_loader = ds.make_loaders(loader, 231 | assign_random_labels, 232 | workers=8, 233 | batch_size=128) 234 | >>> for im, lab in output_loader: 235 | >>> # Do stuff... 236 | ''' 237 | 238 | new_ims = [] 239 | new_targs = [] 240 | total_len = len(loader) 241 | enum_loader = enumerate(loader) 242 | 243 | it = enum_loader if not do_tqdm else tqdm(enum_loader, total=total_len) 244 | for i, (im, targ) in it: 245 | new_im, new_targ = func(im, targ) 246 | if augment or (i / float(total_len) > fraction): 247 | new_ims.append(im.cpu()) 248 | new_targs.append(targ.cpu()) 249 | if i / float(total_len) <= fraction: 250 | new_ims.append(new_im.cpu()) 251 | new_targs.append(new_targ.cpu()) 252 | 253 | dataset = folder.TensorDataset(ch.cat(new_ims, 0), ch.cat(new_targs, 0), transform=transforms) 254 | return ch.utils.data.DataLoader(dataset, num_workers=workers, 255 | batch_size=batch_size, shuffle=shuffle) 256 | -------------------------------------------------------------------------------- /robustness/main.py: -------------------------------------------------------------------------------- 1 | """ 2 | The main file, which exposes the robustness command-line tool, detailed in 3 | :doc:`this walkthrough <../example_usage/cli_usage>`. 4 | """ 5 | 6 | from argparse import ArgumentParser 7 | import os 8 | import git 9 | import torch as ch 10 | 11 | import cox 12 | import cox.utils 13 | import cox.store 14 | 15 | try: 16 | from .model_utils import make_and_restore_model 17 | from .datasets import DATASETS 18 | from .train import train_model, eval_model 19 | from .tools import constants, helpers 20 | from . import defaults, __version__ 21 | from .defaults import check_and_fill_args 22 | except: 23 | raise ValueError("Make sure to run with python -m (see README.md)") 24 | 25 | 26 | parser = ArgumentParser() 27 | parser = defaults.add_args_to_parser(defaults.CONFIG_ARGS, parser) 28 | parser = defaults.add_args_to_parser(defaults.MODEL_LOADER_ARGS, parser) 29 | parser = defaults.add_args_to_parser(defaults.TRAINING_ARGS, parser) 30 | parser = defaults.add_args_to_parser(defaults.PGD_ARGS, parser) 31 | 32 | def main(args, store=None): 33 | '''Given arguments from `setup_args` and a store from `setup_store`, 34 | trains as a model. Check out the argparse object in this file for 35 | argument options. 36 | ''' 37 | # MAKE DATASET AND LOADERS 38 | data_path = os.path.expandvars(args.data) 39 | dataset = DATASETS[args.dataset](data_path) 40 | 41 | train_loader, val_loader = dataset.make_loaders(args.workers, 42 | args.batch_size, data_aug=bool(args.data_aug)) 43 | 44 | train_loader = helpers.DataPrefetcher(train_loader) 45 | val_loader = helpers.DataPrefetcher(val_loader) 46 | loaders = (train_loader, val_loader) 47 | 48 | # MAKE MODEL 49 | model, checkpoint = make_and_restore_model(arch=args.arch, 50 | dataset=dataset, resume_path=args.resume) 51 | if 'module' in dir(model): model = model.module 52 | 53 | print(args) 54 | if args.eval_only: 55 | return eval_model(args, model, val_loader, store=store) 56 | 57 | if not args.resume_optimizer: checkpoint = None 58 | model = train_model(args, model, loaders, store=store, 59 | checkpoint=checkpoint) 60 | return model 61 | 62 | def setup_args(args): 63 | ''' 64 | Fill the args object with reasonable defaults from 65 | :mod:`robustness.defaults`, and also perform a sanity check to make sure no 66 | args are missing. 67 | ''' 68 | # override non-None values with optional config_path 69 | if args.config_path: 70 | args = cox.utils.override_json(args, args.config_path) 71 | 72 | ds_class = DATASETS[args.dataset] 73 | args = check_and_fill_args(args, defaults.CONFIG_ARGS, ds_class) 74 | 75 | if not args.eval_only: 76 | args = check_and_fill_args(args, defaults.TRAINING_ARGS, ds_class) 77 | 78 | if args.adv_train or args.adv_eval: 79 | args = check_and_fill_args(args, defaults.PGD_ARGS, ds_class) 80 | 81 | args = check_and_fill_args(args, defaults.MODEL_LOADER_ARGS, ds_class) 82 | if args.eval_only: assert args.resume is not None, \ 83 | "Must provide a resume path if only evaluating" 84 | return args 85 | 86 | def setup_store_with_metadata(args): 87 | ''' 88 | Sets up a store for training according to the arguments object. See the 89 | argparse object above for options. 90 | ''' 91 | # Add git commit to args 92 | try: 93 | repo = git.Repo(path=os.path.dirname(os.path.realpath(__file__)), 94 | search_parent_directories=True) 95 | version = repo.head.object.hexsha 96 | except git.exc.InvalidGitRepositoryError: 97 | version = __version__ 98 | args.version = version 99 | 100 | # Create the store 101 | store = cox.store.Store(args.out_dir, args.exp_name) 102 | args_dict = args.__dict__ 103 | schema = cox.store.schema_from_dict(args_dict) 104 | store.add_table('metadata', schema) 105 | store['metadata'].append_row(args_dict) 106 | 107 | return store 108 | 109 | if __name__ == "__main__": 110 | args = parser.parse_args() 111 | args = cox.utils.Parameters(args.__dict__) 112 | 113 | args = setup_args(args) 114 | store = setup_store_with_metadata(args) 115 | 116 | final_model = main(args, store=store) 117 | -------------------------------------------------------------------------------- /robustness/model_utils.py: -------------------------------------------------------------------------------- 1 | import torch as ch 2 | from torch import nn 3 | import dill 4 | import os 5 | from .tools import helpers, constants 6 | from .attacker import AttackerModel 7 | 8 | class FeatureExtractor(ch.nn.Module): 9 | ''' 10 | Tool for extracting layers from models. 11 | 12 | Args: 13 | submod (torch.nn.Module): model to extract activations from 14 | layers (list of functions): list of functions where each function, 15 | when applied to submod, returns a desired layer. For example, one 16 | function could be `lambda model: model.layer1`. 17 | 18 | Returns: 19 | A model whose forward function returns the activations from the layers 20 | corresponding to the functions in `layers` (in the order that the 21 | functions were passed in the list). 22 | ''' 23 | def __init__(self, submod, layers): 24 | # layers must be in order 25 | super(FeatureExtractor, self).__init__() 26 | self.submod = submod 27 | self.layers = layers 28 | self.n = 0 29 | 30 | for layer_func in layers: 31 | layer = layer_func(self.submod) 32 | def hook(module, _, output): 33 | module.register_buffer('activations', output) 34 | 35 | layer.register_forward_hook(hook) 36 | 37 | def forward(self, *args, **kwargs): 38 | """ 39 | """ 40 | # self.layer_outputs = {} 41 | out = self.submod(*args, **kwargs) 42 | activs = [layer_fn(self.submod).activations for layer_fn in self.layers] 43 | return [out] + activs 44 | 45 | class DummyModel(nn.Module): 46 | def __init__(self, model): 47 | super().__init__() 48 | self.model = model 49 | 50 | def forward(self, x, *args, **kwargs): 51 | return self.model(x) 52 | 53 | def make_and_restore_model(*_, arch, dataset, resume_path=None, 54 | parallel=False, pytorch_pretrained=False, add_custom_forward=False): 55 | """ 56 | Makes a model and (optionally) restores it from a checkpoint. 57 | 58 | Args: 59 | arch (str|nn.Module): Model architecture identifier or otherwise a 60 | torch.nn.Module instance with the classifier 61 | dataset (Dataset class [see datasets.py]) 62 | resume_path (str): optional path to checkpoint saved with the 63 | robustness library (ignored if ``arch`` is not a string) 64 | not a string 65 | parallel (bool): if True, wrap the model in a DataParallel 66 | (defaults to False) 67 | pytorch_pretrained (bool): if True, try to load a standard-trained 68 | checkpoint from the torchvision library (throw error if failed) 69 | add_custom_forward (bool): ignored unless arch is an instance of 70 | nn.Module (and not a string). Normally, architectures should have a 71 | forward() function which accepts arguments ``with_latent``, 72 | ``fake_relu``, and ``no_relu`` to allow for adversarial manipulation 73 | (see `here` 74 | for more info). If this argument is True, then these options will 75 | not be passed to forward(). (Useful if you just want to train a 76 | model and don't care about these arguments, and are passing in an 77 | arch that you don't want to edit forward() for, e.g. a pretrained model) 78 | Returns: 79 | A tuple consisting of the model (possibly loaded with checkpoint), and the checkpoint itself 80 | """ 81 | if (not isinstance(arch, str)) and add_custom_forward: 82 | arch = DummyModel(arch) 83 | 84 | classifier_model = dataset.get_model(arch, pytorch_pretrained) if \ 85 | isinstance(arch, str) else arch 86 | 87 | model = AttackerModel(classifier_model, dataset) 88 | 89 | # optionally resume from a checkpoint 90 | checkpoint = None 91 | if resume_path and os.path.isfile(resume_path): 92 | print("=> loading checkpoint '{}'".format(resume_path)) 93 | checkpoint = ch.load(resume_path, pickle_module=dill) 94 | 95 | # Makes us able to load models saved with legacy versions 96 | state_dict_path = 'model' 97 | if not ('model' in checkpoint): 98 | state_dict_path = 'state_dict' 99 | 100 | sd = checkpoint[state_dict_path] 101 | sd = {k[len('module.'):]:v for k,v in sd.items()} 102 | model.load_state_dict(sd) 103 | print("=> loaded checkpoint '{}' (epoch {})".format(resume_path, checkpoint['epoch'])) 104 | elif resume_path: 105 | error_msg = "=> no checkpoint found at '{}'".format(resume_path) 106 | raise ValueError(error_msg) 107 | 108 | if parallel: 109 | model = ch.nn.DataParallel(model) 110 | model = model.cuda() 111 | 112 | return model, checkpoint 113 | 114 | def model_dataset_from_store(s, overwrite_params={}, which='last'): 115 | ''' 116 | Given a store directory corresponding to a trained model, return the 117 | original model, dataset object, and args corresponding to the arguments. 118 | ''' 119 | # which options: {'best', 'last', integer} 120 | if type(s) is tuple: 121 | s, e = s 122 | s = cox.store.Store(s, e, mode='r') 123 | 124 | m = s['metadata'] 125 | df = s['metadata'].df 126 | 127 | args = df.to_dict() 128 | args = {k:v[0] for k,v in args.items()} 129 | fns = [lambda x: m.get_object(x), lambda x: m.get_pickle(x)] 130 | conds = [lambda x: m.schema[x] == s.OBJECT, lambda x: m.schema[x] == s.PICKLE] 131 | for fn, cond in zip(fns, conds): 132 | args = {k:(fn(v) if cond(k) else v) for k,v in args.items()} 133 | 134 | args.update(overwrite_params) 135 | args = Parameters(args) 136 | 137 | data_path = os.path.expandvars(args.data) 138 | if not data_path: 139 | data_path = '/tmp/' 140 | 141 | dataset = DATASETS[args.dataset](data_path) 142 | 143 | if which == 'last': 144 | resume = os.path.join(s.path, constants.CKPT_NAME) 145 | elif which == 'best': 146 | resume = os.path.join(s.path, constants.CKPT_NAME_BEST) 147 | else: 148 | assert isinstance(which, int), "'which' must be one of {'best', 'last', int}" 149 | resume = os.path.join(s.path, ckpt_at_epoch(which)) 150 | 151 | model, _ = make_and_restore_model(arch=args.arch, dataset=dataset, 152 | resume_path=resume, parallel=False) 153 | return model, dataset, args 154 | -------------------------------------------------------------------------------- /robustness/tools/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MadryLab/robustness/a9541241defd9972e9334bfcdb804f6aefe24dc7/robustness/tools/__init__.py -------------------------------------------------------------------------------- /robustness/tools/constants.py: -------------------------------------------------------------------------------- 1 | import torch as ch 2 | import cox 3 | from cox import store 4 | 5 | # dog (117), cat (5), frog (3), turtle (5), bird (21), 6 | # monkey (14), fish (9), crab (4), insect (20) 7 | RESTRICTED_IMAGNET_RANGES = [(151, 268), (281, 285), 8 | (30, 32), (33, 37), (80, 100), (365, 382), 9 | (389, 397), (118, 121), (300, 319)] 10 | 11 | CKPT_NAME = 'checkpoint.pt' 12 | BEST_APPEND = '.best' 13 | CKPT_NAME_LATEST = CKPT_NAME + '.latest' 14 | CKPT_NAME_BEST = CKPT_NAME + BEST_APPEND 15 | 16 | ATTACK_KWARG_KEYS = [ 17 | 'criterion', 18 | 'constraint', 19 | 'eps', 20 | 'step_size', 21 | 'iterations', 22 | 'random_start', 23 | 'random_restarts'] 24 | 25 | LOGS_SCHEMA = { 26 | 'epoch':int, 27 | 'nat_prec1':float, 28 | 'adv_prec1':float, 29 | 'nat_loss':float, 30 | 'adv_loss':float, 31 | 'train_prec1':float, 32 | 'train_loss':float, 33 | 'time':float 34 | } 35 | 36 | LOGS_TABLE = 'logs' 37 | 38 | -------------------------------------------------------------------------------- /robustness/tools/custom_modules.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | ch = torch 4 | 5 | class FakeReLU(torch.autograd.Function): 6 | @staticmethod 7 | def forward(ctx, input): 8 | return input.clamp(min=0) 9 | 10 | @staticmethod 11 | def backward(ctx, grad_output): 12 | return grad_output 13 | 14 | class FakeReLUM(nn.Module): 15 | def forward(self, x): 16 | return FakeReLU.apply(x) 17 | 18 | class SequentialWithArgs(torch.nn.Sequential): 19 | def forward(self, input, *args, **kwargs): 20 | vs = list(self._modules.values()) 21 | l = len(vs) 22 | for i in range(l): 23 | if i == l-1: 24 | input = vs[i](input, *args, **kwargs) 25 | else: 26 | input = vs[i](input) 27 | return input 28 | -------------------------------------------------------------------------------- /robustness/tools/folder.py: -------------------------------------------------------------------------------- 1 | import torch.utils.data as data 2 | from torch.utils.data import Dataset 3 | from torchvision import transforms 4 | 5 | from PIL import Image 6 | 7 | import os 8 | import os.path 9 | import sys 10 | 11 | 12 | def has_file_allowed_extension(filename, extensions): 13 | """Checks if a file is an allowed extension. 14 | 15 | Args: 16 | filename (string): path to a file 17 | extensions (iterable of strings): extensions to consider (lowercase) 18 | 19 | Returns: 20 | bool: True if the filename ends with one of given extensions 21 | """ 22 | filename_lower = filename.lower() 23 | return any(filename_lower.endswith(ext) for ext in extensions) 24 | 25 | 26 | def is_image_file(filename): 27 | """Checks if a file is an allowed image extension. 28 | 29 | Args: 30 | filename (string): path to a file 31 | 32 | Returns: 33 | bool: True if the filename ends with a known image extension 34 | """ 35 | return has_file_allowed_extension(filename, IMG_EXTENSIONS) 36 | 37 | 38 | def make_dataset(dir, class_to_idx, extensions): 39 | images = [] 40 | dir = os.path.expanduser(dir) 41 | for target in sorted(class_to_idx.keys()): 42 | d = os.path.join(dir, target) 43 | if not os.path.isdir(d): 44 | continue 45 | 46 | for root, _, fnames in sorted(os.walk(d)): 47 | for fname in sorted(fnames): 48 | if has_file_allowed_extension(fname, extensions): 49 | path = os.path.join(root, fname) 50 | item = (path, class_to_idx[target]) 51 | images.append(item) 52 | 53 | return images 54 | 55 | 56 | class DatasetFolder(data.Dataset): 57 | """A generic data loader where the samples are arranged in this way: :: 58 | 59 | root/class_x/xxx.ext 60 | root/class_x/xxy.ext 61 | root/class_x/xxz.ext 62 | 63 | root/class_y/123.ext 64 | root/class_y/nsdf3.ext 65 | root/class_y/asd932_.ext 66 | 67 | Args: 68 | root (string): Root directory path. 69 | loader (callable): A function to load a sample given its path. 70 | extensions (list[string]): A list of allowed extensions. 71 | transform (callable, optional): A function/transform that takes in 72 | a sample and returns a transformed version. 73 | E.g, ``transforms.RandomCrop`` for images. 74 | target_transform (callable, optional): A function/transform that takes 75 | in the target and transforms it. 76 | 77 | Attributes: 78 | classes (list): List of the class names. 79 | class_to_idx (dict): Dict with items (class_name, class_index). 80 | samples (list): List of (sample path, class_index) tuples 81 | targets (list): The class_index value for each image in the dataset 82 | """ 83 | 84 | def __init__(self, root, loader, extensions, transform=None, 85 | target_transform=None, label_mapping=None): 86 | classes, class_to_idx = self._find_classes(root) 87 | if label_mapping is not None: 88 | classes, class_to_idx = label_mapping(classes, class_to_idx) 89 | 90 | samples = make_dataset(root, class_to_idx, extensions) 91 | if len(samples) == 0: 92 | raise(RuntimeError("Found 0 files in subfolders of: " + root + "\n" 93 | "Supported extensions are: " + ",".join(extensions))) 94 | 95 | self.root = root 96 | self.loader = loader 97 | self.extensions = extensions 98 | 99 | self.classes = classes 100 | self.class_to_idx = class_to_idx 101 | self.samples = samples 102 | self.targets = [s[1] for s in samples] 103 | 104 | self.transform = transform 105 | self.target_transform = target_transform 106 | 107 | def _find_classes(self, dir): 108 | """ 109 | Finds the class folders in a dataset. 110 | 111 | Args: 112 | dir (string): Root directory path. 113 | 114 | Returns: 115 | tuple: (classes, class_to_idx) where classes are relative to (dir), and class_to_idx is a dictionary. 116 | 117 | Ensures: 118 | No class is a subdirectory of another. 119 | """ 120 | if sys.version_info >= (3, 5): 121 | # Faster and available in Python 3.5 and above 122 | classes = [d.name for d in os.scandir(dir) if d.is_dir()] 123 | else: 124 | classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))] 125 | classes.sort() 126 | class_to_idx = {classes[i]: i for i in range(len(classes))} 127 | return classes, class_to_idx 128 | 129 | def __getitem__(self, index): 130 | """ 131 | Args: 132 | index (int): Index 133 | 134 | Returns: 135 | tuple: (sample, target) where target is class_index of the target class. 136 | """ 137 | path, target = self.samples[index] 138 | sample = self.loader(path) 139 | if self.transform is not None: 140 | sample = self.transform(sample) 141 | if self.target_transform is not None: 142 | target = self.target_transform(target) 143 | 144 | return sample, target 145 | 146 | def __len__(self): 147 | return len(self.samples) 148 | 149 | def __repr__(self): 150 | fmt_str = 'Dataset ' + self.__class__.__name__ + '\n' 151 | fmt_str += ' Number of datapoints: {}\n'.format(self.__len__()) 152 | fmt_str += ' Root Location: {}\n'.format(self.root) 153 | tmp = ' Transforms (if any): ' 154 | fmt_str += '{0}{1}\n'.format(tmp, self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp))) 155 | tmp = ' Target Transforms (if any): ' 156 | fmt_str += '{0}{1}'.format(tmp, self.target_transform.__repr__().replace('\n', '\n' + ' ' * len(tmp))) 157 | return fmt_str 158 | 159 | 160 | IMG_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif'] 161 | 162 | 163 | def pil_loader(path): 164 | # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835) 165 | with open(path, 'rb') as f: 166 | img = Image.open(f) 167 | return img.convert('RGB') 168 | 169 | 170 | def accimage_loader(path): 171 | import accimage 172 | try: 173 | return accimage.Image(path) 174 | except IOError: 175 | # Potentially a decoding problem, fall back to PIL.Image 176 | return pil_loader(path) 177 | 178 | 179 | def default_loader(path): 180 | from torchvision import get_image_backend 181 | if get_image_backend() == 'accimage': 182 | return accimage_loader(path) 183 | else: 184 | return pil_loader(path) 185 | 186 | 187 | class ImageFolder(DatasetFolder): 188 | """A generic data loader where the images are arranged in this way: :: 189 | 190 | root/dog/xxx.png 191 | root/dog/xxy.png 192 | root/dog/xxz.png 193 | 194 | root/cat/123.png 195 | root/cat/nsdf3.png 196 | root/cat/asd932_.png 197 | 198 | Args: 199 | root (string): Root directory path. 200 | transform (callable, optional): A function/transform that takes in an PIL image 201 | and returns a transformed version. E.g, ``transforms.RandomCrop`` 202 | target_transform (callable, optional): A function/transform that takes in the 203 | target and transforms it. 204 | loader (callable, optional): A function to load an image given its path. 205 | 206 | Attributes: 207 | classes (list): List of the class names. 208 | class_to_idx (dict): Dict with items (class_name, class_index). 209 | imgs (list): List of (image path, class_index) tuples 210 | """ 211 | def __init__(self, root, transform=None, target_transform=None, 212 | loader=default_loader, label_mapping=None): 213 | super(ImageFolder, self).__init__(root, loader, IMG_EXTENSIONS, 214 | transform=transform, 215 | target_transform=target_transform, 216 | label_mapping=label_mapping) 217 | self.imgs = self.samples 218 | 219 | class TensorDataset(Dataset): 220 | """Dataset wrapping tensors. 221 | 222 | Each sample will be retrieved by indexing tensors along the first dimension. 223 | 224 | Arguments: 225 | *tensors (Tensor): tensors that have the same size of the first dimension. 226 | """ 227 | 228 | def __init__(self, *tensors, transform=None): 229 | assert all(tensors[0].size(0) == tensor.size(0) for tensor in tensors) 230 | self.tensors = tensors 231 | self.transform = transform 232 | 233 | def __getitem__(self, index): 234 | im, targ = tuple(tensor[index] for tensor in self.tensors) 235 | 236 | if self.transform: 237 | real_transform = transforms.Compose([ 238 | transforms.ToPILImage(), 239 | self.transform 240 | ]) 241 | im = real_transform(im) 242 | 243 | return im, targ 244 | 245 | def __len__(self): 246 | return self.tensors[0].size(0) 247 | -------------------------------------------------------------------------------- /robustness/tools/helpers.py: -------------------------------------------------------------------------------- 1 | import torch as ch 2 | 3 | import shutil 4 | import dill 5 | import os 6 | from subprocess import Popen, PIPE 7 | import pandas as pd 8 | from PIL import Image 9 | from . import constants 10 | 11 | def has_attr(obj, k): 12 | """Checks both that obj.k exists and is not equal to None""" 13 | try: 14 | return (getattr(obj, k) is not None) 15 | except KeyError as e: 16 | return False 17 | except AttributeError as e: 18 | return False 19 | 20 | def calc_est_grad(func, x, y, rad, num_samples): 21 | B, *_ = x.shape 22 | Q = num_samples//2 23 | N = len(x.shape) - 1 24 | with ch.no_grad(): 25 | # Q * B * C * H * W 26 | extender = [1]*N 27 | queries = x.repeat(Q, *extender) 28 | noise = ch.randn_like(queries) 29 | norm = noise.view(B*Q, -1).norm(dim=-1).view(B*Q, *extender) 30 | noise = noise / norm 31 | noise = ch.cat([-noise, noise]) 32 | queries = ch.cat([queries, queries]) 33 | y_shape = [1] * (len(y.shape) - 1) 34 | l = func(queries + rad * noise, y.repeat(2*Q, *y_shape)).view(-1, *extender) 35 | grad = (l.view(2*Q, B, *extender) * noise.view(2*Q, B, *noise.shape[1:])).mean(dim=0) 36 | return grad 37 | 38 | def ckpt_at_epoch(num): 39 | return '%s_%s' % (num, constants.CKPT_NAME) 40 | 41 | def accuracy(output, target, topk=(1,), exact=False): 42 | """ 43 | Computes the top-k accuracy for the specified values of k 44 | 45 | Args: 46 | output (ch.tensor) : model output (N, classes) or (N, attributes) 47 | for sigmoid/multitask binary classification 48 | target (ch.tensor) : correct labels (N,) [multiclass] or (N, 49 | attributes) [multitask binary] 50 | topk (tuple) : for each item "k" in this tuple, this method 51 | will return the top-k accuracy 52 | exact (bool) : whether to return aggregate statistics (if 53 | False) or per-example correctness (if True) 54 | 55 | Returns: 56 | A list of top-k accuracies. 57 | """ 58 | with ch.no_grad(): 59 | # Binary Classification 60 | if len(target.shape) > 1: 61 | assert output.shape == target.shape, \ 62 | "Detected binary classification but output shape != target shape" 63 | return [ch.round(ch.sigmoid(output)).eq(ch.round(target)).float().mean()], [-1.0] 64 | 65 | maxk = max(topk) 66 | batch_size = target.size(0) 67 | 68 | _, pred = output.topk(maxk, 1, True, True) 69 | pred = pred.t() 70 | correct = pred.eq(target.view(1, -1).expand_as(pred)) 71 | 72 | res = [] 73 | res_exact = [] 74 | for k in topk: 75 | correct_k = correct[:k].reshape(-1).float() 76 | ck_sum = correct_k.sum(0, keepdim=True) 77 | res.append(ck_sum.mul_(100.0 / batch_size)) 78 | res_exact.append(correct_k) 79 | 80 | if not exact: 81 | return res 82 | else: 83 | return res_exact 84 | 85 | class InputNormalize(ch.nn.Module): 86 | ''' 87 | A module (custom layer) for normalizing the input to have a fixed 88 | mean and standard deviation (user-specified). 89 | ''' 90 | def __init__(self, new_mean, new_std): 91 | super(InputNormalize, self).__init__() 92 | new_std = new_std[..., None, None] 93 | new_mean = new_mean[..., None, None] 94 | 95 | self.register_buffer("new_mean", new_mean) 96 | self.register_buffer("new_std", new_std) 97 | 98 | def forward(self, x): 99 | x = ch.clamp(x, 0, 1) 100 | x_normalized = (x - self.new_mean)/self.new_std 101 | return x_normalized 102 | 103 | class DataPrefetcher(): 104 | def __init__(self, loader, stop_after=None): 105 | self.loader = loader 106 | self.dataset = loader.dataset 107 | self.stream = ch.cuda.Stream() 108 | self.stop_after = stop_after 109 | self.next_input = None 110 | self.next_target = None 111 | 112 | def __len__(self): 113 | return len(self.loader) 114 | 115 | def preload(self): 116 | try: 117 | self.next_input, self.next_target = next(self.loaditer) 118 | except StopIteration: 119 | self.next_input = None 120 | self.next_target = None 121 | return 122 | with ch.cuda.stream(self.stream): 123 | self.next_input = self.next_input.cuda(non_blocking=True) 124 | self.next_target = self.next_target.cuda(non_blocking=True) 125 | 126 | def __iter__(self): 127 | count = 0 128 | self.loaditer = iter(self.loader) 129 | self.preload() 130 | while self.next_input is not None: 131 | ch.cuda.current_stream().wait_stream(self.stream) 132 | input = self.next_input 133 | target = self.next_target 134 | self.preload() 135 | count += 1 136 | yield input, target 137 | if type(self.stop_after) is int and (count > self.stop_after): 138 | break 139 | 140 | class AverageMeter(object): 141 | """Computes and stores the average and current value""" 142 | def __init__(self): 143 | self.reset() 144 | 145 | def reset(self): 146 | self.val = 0 147 | self.avg = 0 148 | self.sum = 0 149 | self.count = 0 150 | 151 | def update(self, val, n=1): 152 | self.val = val 153 | self.sum += val * n 154 | self.count += n 155 | self.avg = self.sum / self.count 156 | 157 | # ImageNet label mappings 158 | def get_label_mapping(dataset_name, ranges): 159 | if dataset_name == 'imagenet': 160 | label_mapping = None 161 | elif dataset_name == 'restricted_imagenet': 162 | def label_mapping(classes, class_to_idx): 163 | return restricted_label_mapping(classes, class_to_idx, ranges=ranges) 164 | elif dataset_name == 'custom_imagenet': 165 | def label_mapping(classes, class_to_idx): 166 | return custom_label_mapping(classes, class_to_idx, ranges=ranges) 167 | else: 168 | raise ValueError('No such dataset_name %s' % dataset_name) 169 | 170 | return label_mapping 171 | 172 | def restricted_label_mapping(classes, class_to_idx, ranges): 173 | range_sets = [ 174 | set(range(s, e+1)) for s,e in ranges 175 | ] 176 | 177 | # add wildcard 178 | # range_sets.append(set(range(0, 1002))) 179 | mapping = {} 180 | for class_name, idx in class_to_idx.items(): 181 | for new_idx, range_set in enumerate(range_sets): 182 | if idx in range_set: 183 | mapping[class_name] = new_idx 184 | # assert class_name in mapping 185 | filtered_classes = list(mapping.keys()).sort() 186 | return filtered_classes, mapping 187 | 188 | def custom_label_mapping(classes, class_to_idx, ranges): 189 | 190 | mapping = {} 191 | for class_name, idx in class_to_idx.items(): 192 | for new_idx, range_set in enumerate(ranges): 193 | if idx in range_set: 194 | mapping[class_name] = new_idx 195 | 196 | filtered_classes = list(mapping.keys()).sort() 197 | return filtered_classes, mapping 198 | -------------------------------------------------------------------------------- /robustness/tools/openimgs_helpers.py: -------------------------------------------------------------------------------- 1 | import os 2 | import csv 3 | import numpy as np 4 | import torch as ch 5 | import torch.utils.data as data 6 | import robustness.data_augmentation as da 7 | from robustness import imagenet_models 8 | from .folder import default_loader, IMG_EXTENSIONS 9 | 10 | target_transform_oi = ch.Tensor 11 | 12 | def load_class_desc(data_dir): 13 | """Returns map from cid to class name.""" 14 | 15 | class_names = {} 16 | 17 | with open(os.path.join(data_dir, "metadata", "class-descriptions-boxable.csv"), newline="") as csvfile: 18 | for ri, row in enumerate(csv.reader(csvfile, delimiter=' ', quotechar='|')): 19 | cid = row[0].split(',')[0] 20 | cname = ' '.join([row[0].split(',')[1]] + row[1:]) 21 | assert cid not in class_names 22 | class_names[cid] = cname 23 | 24 | return class_names 25 | 26 | def get_image_annotations_mode(class_names, data_dir, mode="train"): 27 | """Returns map from img number to label (along with verification 28 | source and confidence)""" 29 | 30 | assert mode in set(["train", "test", "validation"]) 31 | lab_dir = os.path.join(data_dir, 32 | "labels", 33 | f"{mode}-annotations-human-imagelabels-boxable.csv") 34 | prefix = "oidv6-" if mode == "train" else "" 35 | anno_dir = os.path.join(data_dir, 36 | "boxes", 37 | f"{prefix}{mode}-annotations-bbox.csv") 38 | 39 | img_to_label = {} 40 | with open(lab_dir, newline="") as csvfile: 41 | for ri, row in enumerate(csv.reader(csvfile, delimiter=' ', quotechar='|')): 42 | if ri == 0: continue 43 | 44 | assert len(row) == 1 45 | im_id, ver, cno, conf = tuple(row[0].split(",")) 46 | cno = class_names[cno] 47 | 48 | if im_id not in img_to_label: 49 | img_to_label[im_id] = {} 50 | 51 | if cno not in img_to_label[im_id]: 52 | img_to_label[im_id][cno] = {'ver': [], 'conf': []} 53 | img_to_label[im_id][cno]['ver'].append(ver) 54 | img_to_label[im_id][cno]['conf'].append(conf) 55 | 56 | 57 | for im_id in img_to_label: 58 | for lab in img_to_label[im_id]: 59 | assert len(np.unique(img_to_label[im_id][lab]['conf'])) == 1 60 | img_to_label[im_id][lab]['conf'] = img_to_label[im_id][lab]['conf'][0] 61 | 62 | with open(anno_dir, newline="") as csvfile: 63 | for ri, row in enumerate(csv.reader(csvfile, delimiter=' ', quotechar='|')): 64 | if ri == 0: continue 65 | assert len(row) == 1 66 | rs = row[0].split(",") 67 | im_id, src, cno = tuple(rs[:3]) 68 | cno = class_names[cno] 69 | 70 | box = [float(v) for v in rs[4:8]] 71 | if 'box' not in img_to_label[im_id][cno] or src == 'activemil': 72 | img_to_label[im_id][cno]['box'] = box 73 | 74 | return img_to_label 75 | 76 | 77 | def make_dataset(dir, mode, sample_info, 78 | class_to_idx, class_to_idx_comp, extensions): 79 | 80 | images = [] 81 | allowed_labels = set(class_to_idx.keys()) 82 | Nclasses = len(set(class_to_idx.values())) 83 | 84 | for k, v in sample_info.items(): 85 | 86 | img_path = os.path.join(dir, "images", mode, k + ".jpg") 87 | 88 | pos_labels = set([l for l in v.keys() if v[l]['conf'][0] == '1']) 89 | neg_labels = set([l for l in v.keys() if v[l]['conf'][0] == '0']) 90 | 91 | pos_labels = pos_labels.intersection(allowed_labels) 92 | neg_labels = neg_labels.intersection(allowed_labels) 93 | if Nclasses == 601 or len(pos_labels) != 0: 94 | label = [0] * Nclasses 95 | all_labels = [0] * 601 96 | for p in pos_labels: 97 | label[class_to_idx[p]] = 1 98 | all_labels[class_to_idx_comp[p]] = 1 99 | for n in neg_labels: 100 | if label[class_to_idx[n]] == 0: 101 | label[class_to_idx[n]] = -1 102 | all_labels[class_to_idx_comp[n]] = -1 103 | item = (img_path, label, all_labels) 104 | images.append(item) 105 | 106 | return images 107 | 108 | class OIDatasetFolder(data.Dataset): 109 | """A generic data loader where the samples are arranged in this way: :: 110 | Args: 111 | root (string): Root directory path. 112 | loader (callable): A function to load a sample given its path. 113 | extensions (list[string]): A list of allowed extensions. 114 | transform (callable, optional): A function/transform that takes in 115 | a sample and returns a transformed version. 116 | E.g, ``transforms.RandomCrop`` for images. 117 | target_transform (callable, optional): A function/transform that takes 118 | in the target and transforms it. 119 | Attributes: 120 | classes (list): List of the class names. 121 | class_to_idx (dict): Dict with items (class_name, class_index). 122 | samples (list): List of (sample path, class_index) tuples 123 | targets (list): The class_index value for each image in the dataset 124 | """ 125 | 126 | def __init__(self, root, train=True, extensions=IMG_EXTENSIONS, 127 | loader=default_loader, transform=None, 128 | target_transform=target_transform_oi, label_mapping=None, 129 | download=False): 130 | classes, class_to_idx, code_to_class = self._find_classes(root) 131 | class_to_idx_comp = {k: v for k, v in class_to_idx.items()} 132 | if label_mapping is not None: 133 | classes, class_to_idx = label_mapping(classes, class_to_idx) 134 | 135 | mode = "train" if train else "test" 136 | sample_info = get_image_annotations_mode(code_to_class, 137 | mode=mode, 138 | data_dir=root) 139 | 140 | 141 | samples = make_dataset(root, mode, sample_info, 142 | class_to_idx, class_to_idx_comp, 143 | extensions) 144 | if len(samples) == 0: 145 | raise(RuntimeError("Found 0 files in subfolders of: " + root + "\n" 146 | "Supported extensions are: " + ",".join(extensions))) 147 | 148 | self.root = root 149 | self.loader = loader 150 | self.extensions = extensions 151 | 152 | self.classes = classes 153 | self.class_to_idx = class_to_idx 154 | self.samples = samples 155 | self.targets = [s[1] for s in samples] 156 | self.all_targets = [s[2] for s in samples] 157 | 158 | self.transform = transform 159 | self.target_transform = target_transform 160 | 161 | def _find_classes(self, dir): 162 | """ 163 | Finds the class folders in a dataset. 164 | """ 165 | code_to_class = load_class_desc(dir) 166 | classes = [v for v in code_to_class.values()] 167 | class_to_idx = {code_to_class[k]: i for i, k in enumerate(code_to_class)} 168 | return classes, class_to_idx, code_to_class 169 | 170 | def __getitem__(self, index): 171 | """ 172 | Args: 173 | index (int): Index 174 | Returns: 175 | tuple: (sample, target) where target is class_index of the target class. 176 | """ 177 | path, target, comp_target = self.samples[index] 178 | sample = self.loader(path) 179 | if self.transform is not None: 180 | sample = self.transform(sample) 181 | if self.target_transform is not None: 182 | target = self.target_transform(target) 183 | if self.target_transform is not None: 184 | comp_target = self.target_transform(comp_target) 185 | return sample, target, comp_target 186 | 187 | def __len__(self): 188 | return len(self.samples) 189 | 190 | def __repr__(self): 191 | fmt_str = 'Dataset ' + self.__class__.__name__ + '\n' 192 | fmt_str += ' Number of datapoints: {}\n'.format(self.__len__()) 193 | fmt_str += ' Root Location: {}\n'.format(self.root) 194 | tmp = ' Transforms (if any): ' 195 | fmt_str += '{0}{1}\n'.format(tmp, self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp))) 196 | tmp = ' Target Transforms (if any): ' 197 | fmt_str += '{0}{1}'.format(tmp, self.target_transform.__repr__().replace('\n', '\n' + ' ' * len(tmp))) 198 | return fmt_str 199 | 200 | def get_label_map(data_dir): 201 | CLASS_NAMES = load_class_desc(data_dir) 202 | label_map = {i: v for i, v in enumerate(CLASS_NAMES.values())} 203 | return label_map 204 | 205 | def get_labels(targ, label_map): 206 | pos_labels, neg_labels = [], [] 207 | for ti, t in enumerate(targ.numpy()): 208 | if t == 1: 209 | pos_labels.append(f"+ {label_map[ti]}") 210 | elif t == -1: 211 | neg_labels.append(f"- {label_map[ti]}") 212 | return ", ".join(pos_labels) + " | " + ", ".join(neg_labels) 213 | -------------------------------------------------------------------------------- /robustness/tools/vis_tools.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import matplotlib.pyplot as plt 3 | from sklearn.decomposition import PCA 4 | from sklearn import manifold 5 | import seaborn as sns 6 | 7 | def get_axis(axarr, H, W, i, j): 8 | H, W = H - 1, W - 1 9 | if not (H or W): 10 | ax = axarr 11 | elif not (H and W): 12 | ax = axarr[max(i, j)] 13 | else: 14 | ax = axarr[i][j] 15 | return ax 16 | 17 | def show_image_row(xlist, ylist=None, fontsize=12, size=(2.5, 2.5), tlist=None, filename=None): 18 | H, W = len(xlist), len(xlist[0]) 19 | fig, axarr = plt.subplots(H, W, figsize=(size[0] * W, size[1] * H)) 20 | for w in range(W): 21 | for h in range(H): 22 | ax = get_axis(axarr, H, W, h, w) 23 | ax.imshow(xlist[h][w].permute(1, 2, 0)) 24 | ax.xaxis.set_ticks([]) 25 | ax.yaxis.set_ticks([]) 26 | ax.xaxis.set_ticklabels([]) 27 | ax.yaxis.set_ticklabels([]) 28 | if ylist and w == 0: 29 | ax.set_ylabel(ylist[h], fontsize=fontsize) 30 | if tlist: 31 | ax.set_title(tlist[h][w], fontsize=fontsize) 32 | if filename is not None: 33 | plt.savefig(filename, bbox_inches='tight') 34 | plt.show() 35 | 36 | 37 | def show_image_column(xlist, ylist=None, fontsize=12, size=(2.5, 2.5), tlist=None, filename=None): 38 | W, H = len(xlist), len(xlist[0]) 39 | fig, axarr = plt.subplots(H, W, figsize=(size[0] * W, size[1] * H)) 40 | for w in range(W): 41 | for h in range(H): 42 | ax = get_axis(axarr, H, W, h, w) 43 | ax.imshow(xlist[w][h].permute(1, 2, 0)) 44 | ax.xaxis.set_ticks([]) 45 | ax.yaxis.set_ticks([]) 46 | ax.xaxis.set_ticklabels([]) 47 | ax.yaxis.set_ticklabels([]) 48 | if ylist and h == 0: 49 | ax.set_title(ylist[w], fontsize=fontsize) 50 | if tlist: 51 | ax.set_title(tlist[w][h], fontsize=fontsize) 52 | if filename is not None: 53 | plt.savefig(filename, bbox_inches='tight') 54 | plt.show() 55 | 56 | def filter_data(metadata, criteria, value): 57 | crit = [True] * len(metadata) 58 | for c, v in zip(criteria, value): 59 | v = [v] if not isinstance(v, list) else v 60 | crit &= metadata[c].isin(v) 61 | metadata_int = metadata[crit] 62 | exp_ids = metadata_int['exp_id'].tolist() 63 | return exp_ids 64 | 65 | def plot_axis(ax, x, y, xlabel, ylabel, **kwargs): 66 | ax.plot(x, y, **kwargs) 67 | ax.set_xlabel(xlabel, fontsize=14) 68 | ax.set_ylabel(ylabel, fontsize=14) 69 | 70 | 71 | def plot_tsne(x, y, npca=50, markersize=10): 72 | Xlow = PCA(n_components=npca).fit_transform(x) 73 | Y = manifold.TSNE(n_components=2).fit_transform(Xlow) 74 | palette = sns.color_palette("Paired", len(np.unique(y))) 75 | color_dict = {l: c for l, c in zip(range(len(np.unique(y))), palette)} 76 | colors = [color_dict[l] for l in y] 77 | plt.scatter(Y[:, 0], Y[:, 1], markersize, colors, 'o') 78 | plt.show() 79 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | """A setuptools based setup module. 2 | See: 3 | https://packaging.python.org/en/latest/distributing.html 4 | https://github.com/pypa/sampleproject 5 | """ 6 | 7 | # Always prefer setuptools over distutils 8 | from setuptools import setup, find_packages 9 | # To use a consistent encoding 10 | from codecs import open 11 | from os import path 12 | 13 | here = path.abspath(path.dirname(__file__)) 14 | 15 | # Get the long description from the README file 16 | with open(path.join(here, 'README.rst'), encoding='utf-8') as f: 17 | long_description = f.read() 18 | 19 | setup( 20 | name='robustness', 21 | 22 | # Versions should comply with PEP440. For a discussion on single-sourcing 23 | # the version across setup.py and the project code, see 24 | # https://packaging.python.org/en/latest/single_source_version.html 25 | version='1.2.1.post2', 26 | 27 | description='Tools for Robustness', 28 | long_description=long_description, 29 | long_description_content_type='text/x-rst', 30 | 31 | # The project's main homepage. 32 | #url='https://github.com/', 33 | 34 | # Author details 35 | author='MadryLab', 36 | author_email='ailyas@mit.edu', 37 | 38 | # Choose your license 39 | license='MIT', 40 | 41 | # See https://pypi.python.org/pypi?%4Aaction=list_classifiers 42 | classifiers=[ 43 | # How mature is this project? Common values are 44 | # 3 - Alpha 45 | # 4 - Beta 46 | # 5 - Production/Stable 47 | 'Development Status :: 4 - Beta', 48 | 49 | # Indicate who your project is intended for 50 | 'Intended Audience :: Developers', 51 | 'Topic :: Software Development :: Build Tools', 52 | 53 | # Pick your license as you wish (should match "license" above) 54 | 'License :: OSI Approved :: MIT License', 55 | 56 | # Specify the Python versions you support here. In particular, ensure 57 | # that you indicate whether you support Python 2, Python 3 or both. 58 | 'Programming Language :: Python :: 3', 59 | 'Programming Language :: Python :: 3.5', 60 | 'Programming Language :: Python :: 3.6', 61 | ], 62 | 63 | # What does your project relate to? 64 | keywords='logging tools madrylab', 65 | 66 | # You can just specify the packages manually here if your project is 67 | # simple. Or you can use find_packages(). 68 | packages=['robustness', 69 | 'robustness.tools', 70 | 'robustness.cifar_models', 71 | 'robustness.imagenet_models' 72 | ], 73 | 74 | include_package_data=True, 75 | package_data={ 76 | 'certificate': ['client/server.crt'] 77 | }, 78 | 79 | # Alternatively, if you want to distribute just a my_module.py, uncomment 80 | # this: 81 | # py_modules=["my_module"], 82 | # 83 | # List run-time dependencies here. These will be installed by pip when 84 | # your project is installed. For an analysis of "install_requires" vs pip's 85 | # requirements files see: 86 | # https://packaging.python.org/en/latest/requirements.html 87 | install_requires=['tqdm', 'grpcio', 'psutil', 'gitpython','py3nvml', 'cox', 88 | 'scikit-learn', 'seaborn', 'torch', 'torchvision', 'pandas', 89 | 'numpy', 'scipy', 'GPUtil', 'dill', 'tensorboardX', 'tables', 90 | 'matplotlib'], 91 | test_suite='nose.collector', 92 | tests_require=['nose'], 93 | ) 94 | --------------------------------------------------------------------------------