├── .gitattributes
├── .gitignore
├── LICENSE
├── README.md
├── dissect_all.sh
├── experiment
├── __init__.py
├── dissect_experiment.py
├── generator_int_experiment.py
├── intervention_experiment.py
├── oldalexnet.py
├── oldresnet152.py
├── oldvgg16.py
├── proggan.py
├── readdissect.py
├── setting.py
└── shapebias_experiment.py
├── g_intervention.sh
├── intervention.sh
├── netdissect
├── __init__.py
├── bargraph.py
├── easydict.py
├── imgsave.py
├── imgviz.py
├── labwidget.py
├── nethook.py
├── paintwidget.py
├── parallelfolder.py
├── pbar.py
├── pidfile.py
├── renormalize.py
├── report.html
├── runningstats.py
├── sampler.py
├── segmenter.py
├── segmodel
│ ├── __init__.py
│ ├── colors150.npy
│ ├── mobilenet.py
│ ├── models.py
│ ├── object150_info.csv
│ ├── resnet.py
│ └── resnext.py
├── segviz.py
├── show.py
├── tally.py
├── upsample.py
├── upsegmodel
│ ├── __init__.py
│ ├── models.py
│ ├── prroi_pool
│ │ ├── .gitignore
│ │ ├── README.md
│ │ ├── __init__.py
│ │ ├── build.py
│ │ ├── functional.py
│ │ ├── prroi_pool.py
│ │ ├── src
│ │ │ ├── prroi_pooling_gpu.c
│ │ │ ├── prroi_pooling_gpu.h
│ │ │ ├── prroi_pooling_gpu_impl.cu
│ │ │ └── prroi_pooling_gpu_impl.cuh
│ │ └── test_prroi_pooling2d.py
│ ├── resnet.py
│ └── resnext.py
├── workerpool.py
└── zdataset.py
├── notebooks
├── adv_experiment_bedroom.ipynb
├── adv_experiment_plot.ipynb
├── dissect_classifier_experiment.ipynb
├── dissect_generator_experiment.ipynb
├── intervention-classifier-experiment.ipynb
├── intervention-generator-experiment.ipynb
├── ipynb_drop_output.py
├── setup_notebooks.sh
├── shapebias_experiment.ipynb
├── single-classifier-unit-plot.ipynb
└── single-generator-unit-plot.ipynb
├── setup
├── denv.yml
└── setup_denv.sh
├── stylization
├── LICENSE
├── README.md
├── function.py
├── models
│ └── download_models.sh
├── net.py
├── parallel-make-sp.sh
├── stylize.py
└── torch_to_pytorch.py
└── www
├── arxiv-thumb.png
├── classifier-dissection.png
├── classifier-intervention.png
├── dissection-compare.png
├── gandissect-tutorial.png
├── generator-dissection.png
├── generator-intervention.png
├── netdissect-tutorial.png
├── netdissect_code.png
├── paper-thumb.png
├── si-thumb.png
└── website-thumb.png
/.gitattributes:
--------------------------------------------------------------------------------
1 | *.ipynb filter=clean_ipynb
2 |
--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | *.pyc
2 | /datasets
3 | /results
4 | /notebooks/datasets
5 | /notebooks/results
6 | /notebooks/netdissect
7 | /notebooks/experiment
8 | /notebooks/unused
9 | /stylization/data
10 | /stylization/models/*.pth
11 | __pycache__
12 | .DS_Store
13 | .__*
14 | .ipynb*
15 | .nfs*
16 | .*swp
17 | .idea
18 |
--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
1 | Most files in this directory (stylization/) are either directly copied from the
2 | pytorch-AdaIN repository (https://github.com/naoto0804/pytorch-AdaIN)
3 | or adapted slightly. The following license applies to these files:
4 |
5 | MIT License
6 |
7 | Copyright (c) 2018 Naoto Inoue
8 |
9 | Permission is hereby granted, free of charge, to any person obtaining a copy
10 | of this software and associated documentation files (the "Software"), to deal
11 | in the Software without restriction, including without limitation the rights
12 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
13 | copies of the Software, and to permit persons to whom the Software is
14 | furnished to do so, subject to the following conditions:
15 |
16 | The above copyright notice and this permission notice shall be included in all
17 | copies or substantial portions of the Software.
18 |
19 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
20 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
21 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
22 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
23 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
24 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
25 | SOFTWARE.
26 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # What is the Role of a Neuron?
2 |
3 | When a deep network is trained on a high-level task such as classifying a place or synthesizing a scene, individual neural units within the network will often emerge that match specific human-interpretable concepts, like "trees", "windows", or "human faces."
4 |
5 | What role do such individual units serve within a deep network?
6 |
7 | We examine this question in two types of networks that contain interpretable units: networks trained to classify images of scenes (supervised image classifiers), and networks trained to synthesize images of scenes (generative adversarial networks).
8 |
9 | [**Understanding the Role of Individual Units in a Deep Network**](https://dissect.csail.mit.edu/).
10 | [David Bau](https://people.csail.mit.edu/davidbau/home/), [Jun-Yan Zhu](https://www.cs.cmu.edu/~junyanz/), [Hendrik Strobelt](http://hendrik.strobelt.com/), [Agata Lapedriza](https://www.media.mit.edu/people/agata/overview/), [Bolei Zhou](http://bzhou.ie.cuhk.edu.hk/), [Antonio Torralba](http://web.mit.edu/torralba/www/).
11 | Proceedings of the National Academy of Sciences, September 2020.
12 | MIT, MIT-IBM Watson AI Lab, IBM Research, The Chinese University of Hong Kong, Adobe Research
13 |
14 |
',
77 | ]
78 | needs_end = False
79 | table_mode = False
80 | for i, line in enumerate(obj):
81 | if i == 0:
82 | needs_end = True
83 | if isinstance(line, tuple):
84 | table_mode = True
85 | results.append(tstart)
86 | else:
87 | results.append(blockstart)
88 | if table_mode:
89 | results.append(rstart)
90 | if not isinstance(line, str) and hasattr(line, '__iter__'):
91 | for cell in line:
92 | results.append(cstart)
93 | results.extend(blocks_tags(cell))
94 | results.append(cend)
95 | else:
96 | results.append(cstart)
97 | results.extend(blocks_tags(line))
98 | results.append(cend)
99 | results.append(rend)
100 | else:
101 | results.extend(blocks_tags(line))
102 | if needs_end:
103 | results.append(table_mode and tend or blockend)
104 | return results
105 |
106 |
107 | def pil_to_b64(img, format='png'):
108 | buffered = io.BytesIO()
109 | img.save(buffered, format=format)
110 | return base64.b64encode(buffered.getvalue()).decode('utf-8')
111 |
112 |
113 | def pil_to_url(img, format='png'):
114 | return 'data:image/%s;base64,%s' % (format, pil_to_b64(img, format))
115 |
116 |
117 | def pil_to_html(img, margin=1):
118 | mattr = ' style="margin:%dpx"' % margin
119 | return '' % (pil_to_url(img), mattr)
120 |
121 |
122 | def a(x, cols=None):
123 | global g_buffer
124 | if g_buffer is None:
125 | g_buffer = []
126 | g_buffer.append(x)
127 | if cols is not None and len(g_buffer) >= cols:
128 | flush()
129 |
130 |
131 | def reset():
132 | global g_buffer
133 | g_buffer = None
134 |
135 |
136 | def flush(*args, **kwargs):
137 | global g_buffer
138 | if g_buffer is not None:
139 | x = g_buffer
140 | g_buffer = None
141 | display(blocks(x, *args, **kwargs))
142 |
143 |
144 | def show(x=None, *args, **kwargs):
145 | flush(*args, **kwargs)
146 | if x is not None:
147 | display(blocks(x, *args, **kwargs))
148 |
149 |
150 | def html(obj, space=''):
151 | return blocks(obj, space)._repr_html_()
152 |
153 |
154 | class CallableModule(types.ModuleType):
155 | def __init__(self):
156 | # or super().__init__(__name__) for Python 3
157 | types.ModuleType.__init__(self, __name__)
158 | self.__dict__.update(sys.modules[__name__].__dict__)
159 |
160 | def __call__(self, x=None, *args, **kwargs):
161 | show(x, *args, **kwargs)
162 |
163 |
164 | sys.modules[__name__] = CallableModule()
165 |
--------------------------------------------------------------------------------
/netdissect/upsample.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from torchvision import transforms
3 |
4 |
5 | def upsampler(target_shape, data_shape=None,
6 | image_size=None, scale_offset=None,
7 | source=None, convolutions=None, dtype=torch.float, device=None):
8 | '''
9 | Returns a function that will upsample a batch of torch data from the
10 | expected data_shape to the specified target_shape. Can use scale_offset
11 | and image_size to center the grid in a nondefault way: scale_offset
12 | maps feature pixels to image_size pixels, and it is assumed that
13 | the target_shape is a uniform downsampling of image_size.
14 | '''
15 | if source is not None:
16 | assert image_size is None
17 | image_size = image_size_from_source(source)
18 | if convolutions is not None:
19 | assert scale_offset is None
20 | scale_offset = sequence_scale_offset(convolutions)
21 | if image_size is not None and data_shape is None:
22 | data_shape = sequence_data_size(convolutions, image_size)
23 | assert data_shape is not None
24 | assert len(data_shape) == 2
25 | grid = upsample_grid(data_shape, target_shape, image_size, scale_offset,
26 | dtype, device)
27 | batch_grid = grid
28 | # padding mode could be 'border'
29 |
30 | def upsample_func(data, mode='bilinear', padding_mode='zeros'):
31 | nonlocal grid, batch_grid
32 | # Use the same grid over the whole batch
33 | if batch_grid.shape[0] != data.shape[0]:
34 | batch_grid = grid.expand((data.shape[0],) + grid.shape[1:])
35 | if batch_grid.device != data.device:
36 | batch_grid = batch_grid.to(data.device)
37 | try:
38 | return torch.nn.functional.grid_sample(data, batch_grid, mode=mode,
39 | padding_mode=padding_mode, align_corners=True)
40 | except:
41 | return torch.nn.functional.grid_sample(data, batch_grid, mode=mode,
42 | padding_mode=padding_mode) # older pytorch version
43 | return upsample_func
44 |
45 |
46 | def sequence_scale_offset(modulelist):
47 | '''Returns (yscale, yoffset), (xscale, xoffset) given a list of modules.
48 | To convert output coordinates back to input coordinates while preserving
49 | centers of receptive fields, the affine transformation is:
50 | inpx = outx * xscale + xoffset
51 | inpy = outy * yscale + yoffset
52 | In both coordinate systems, (0, 0) refers to the upper-left corner
53 | of the first pixel, (0.5, 0.5) refers to the center of that pixel,
54 | and (1, 1) refers to the lower-right corner of that same pixel.
55 |
56 | Modern convnets tend to add padding to keep receptive fields centered
57 | while scaling, which will result in zero offsets. For example, after resnet
58 | does five stride-2 reductions, the scale_offset is just ((32, 0), (32, 0)).
59 | However, AlexNet does not pad every layer, and after five stride-2
60 | reductions, the scale_offset is ((32, 31), (32, 31)).
61 | '''
62 | return tuple(convconfig_scale_offset(d) for d in convconfigs(modulelist))
63 |
64 |
65 | def sequence_data_size(modulelist, input_size):
66 | '''Returns (yscale, yoffset), (xscale, xoffset) given a list of modules.
67 | To convert output coordinates back to input coordinates while preserving
68 | centers of receptive fields, the affine transformation is:
69 | inpx = outx * xscale + xoffset
70 | inpy = outy * yscale + yoffset
71 | In both coordinate systems, (0, 0) refers to the upper-left corner
72 | of the first pixel, (0.5, 0.5) refers to the center of that pixel,
73 | and (1, 1) refers to the lower-right corner of that same pixel.
74 |
75 | Modern convnets tend to add padding to keep receptive fields centered
76 | while scaling, which will result in zero offsets. For example, after resnet
77 | does five stride-2 reductions, the scale_offset is just ((32, 0), (32, 0)).
78 | However, AlexNet does not pad every layer, and after five stride-2
79 | reductions, the scale_offset is ((32, 31), (32, 31)).
80 | '''
81 | return tuple(convconfig_data_size(d, s)
82 | for d, s in zip(convconfigs(modulelist), input_size))
83 |
84 |
85 | def convconfig_scale_offset(convconfigs):
86 | '''Composes a lists of [(k, d, s, p)...] into a single total scale and
87 | offset that returns to the input coordinates.
88 | '''
89 | if len(convconfigs) == 0:
90 | return (1, 0)
91 | scale, offset = convconfig_scale_offset(convconfigs[1:])
92 | kernel, dilation, stride, padding = convconfigs[0]
93 | scale *= stride
94 | offset *= stride
95 | offset += (kernel - 1) * dilation / 2.0 - padding
96 | return scale, offset
97 |
98 |
99 | def convconfig_data_size(convconfigs, data_size):
100 | '''Applies a list of [(k, d, s, p)...] to the given input size to obtain
101 | an output size.
102 | '''
103 | for kernel, dilation, stride, padding in convconfigs:
104 | data_size = (1 + (data_size + 2 * padding
105 | - dilation * (kernel - 1) - 1) // stride)
106 | return data_size
107 |
108 |
109 | def convconfigs(modulelist):
110 | '''Converts a list of modules to a pair of lists of
111 | [(kernel_size, dilation, stride, padding)...]: one for x, and one for y.'''
112 | result = []
113 | for module in modulelist:
114 | settings = tuple(getattr(module, n, d)
115 | for n, d in (('kernel_size', 1),
116 | ('dilation', 1), ('stride', 1), ('padding', 0)))
117 | settings = tuple((s if isinstance(s, tuple) else (s, s))
118 | for s in settings)
119 | if settings != ((1, 1), (1, 1), (1, 1), (0, 0)):
120 | result.append(zip(*settings))
121 | return list(zip(*result))
122 |
123 |
124 | def upsample_grid(data_shape, target_shape, image_size=None,
125 | scale_offset=None, dtype=torch.float, device=None):
126 | '''Prepares a grid to use with grid_sample to upsample a batch of
127 | features in data_shape to the target_shape. Can use scale_offset
128 | and image_size to center the grid in a nondefault way: scale_offset
129 | maps feature pixels to image_size pixels, and it is assumed that
130 | the target_shape is a uniform downsampling of image_size.'''
131 | # Default is that nothing is resized.
132 | if target_shape is None:
133 | target_shape = data_shape
134 | # Make a default scale_offset to fill the image if there isn't one
135 | if scale_offset is None:
136 | scale = tuple(float(ts) / ds
137 | for ts, ds in zip(target_shape, data_shape))
138 | offset = tuple(0.5 * s - 0.5 for s in scale)
139 | else:
140 | scale, offset = (v for v in zip(*scale_offset))
141 | # Handle downsampling for different input vs target shape.
142 | if image_size is not None:
143 | scale = tuple(s * (ts - 1) / (ns - 1)
144 | for s, ns, ts in zip(scale, image_size, target_shape))
145 | offset = tuple(o * (ts - 1) / (ns - 1)
146 | for o, ns, ts in zip(offset, image_size, target_shape))
147 | # Pytorch needs target coordinates in terms of source coordinates [-1..1]
148 | ty, tx = (((torch.arange(ts, dtype=dtype, device=device) - o)
149 | * (2 / (s * max(1, (ss - 1)))) - 1)
150 | for ts, ss, s, o, in zip(target_shape, data_shape, scale, offset))
151 | # Whoa, note that grid_sample reverses the order y, x -> x, y.
152 | grid = torch.stack(
153 | (tx[None, :].expand(target_shape), ty[:, None].expand(target_shape)), 2
154 | )[None, :, :, :].expand((1, target_shape[0], target_shape[1], 2))
155 | return grid
156 |
157 |
158 | def image_size_from_source(source):
159 | sizer = find_sizer(source)
160 | if sizer is not None:
161 | size = sizer.size
162 | elif hasattr(source, 'resolution'):
163 | size = source.resolution
164 | if hasattr(size, '__len__'):
165 | return size
166 | return (size, size)
167 |
168 |
169 | def find_sizer(source):
170 | '''
171 | Crawl around the transforms attached to a dataset looking for
172 | the last crop or resize transform to return.
173 | '''
174 | if source is None:
175 | return None
176 | if isinstance(source, (transforms.Resize, transforms.RandomCrop,
177 | transforms.RandomResizedCrop, transforms.CenterCrop)):
178 | return source
179 | t = getattr(source, 'transform', None)
180 | if t is not None:
181 | return find_sizer(t)
182 | ts = getattr(source, 'transforms', None)
183 | if ts is not None:
184 | for t in reversed(ts):
185 | result = find_sizer(t)
186 | if result is not None:
187 | return result
188 | return None
189 |
--------------------------------------------------------------------------------
/netdissect/upsegmodel/__init__.py:
--------------------------------------------------------------------------------
1 | from .models import ModelBuilder, SegmentationModule
2 |
--------------------------------------------------------------------------------
/netdissect/upsegmodel/prroi_pool/.gitignore:
--------------------------------------------------------------------------------
1 | *.o
2 | /_prroi_pooling
3 |
--------------------------------------------------------------------------------
/netdissect/upsegmodel/prroi_pool/README.md:
--------------------------------------------------------------------------------
1 | # PreciseRoIPooling
2 | This repo implements the **Precise RoI Pooling** (PrRoI Pooling), proposed in the paper **Acquisition of Localization Confidence for Accurate Object Detection** published at ECCV 2018 (Oral Presentation).
3 |
4 | **Acquisition of Localization Confidence for Accurate Object Detection**
5 |
6 | _Borui Jiang*, Ruixuan Luo*, Jiayuan Mao*, Tete Xiao, Yuning Jiang_ (* indicates equal contribution.)
7 |
8 | https://arxiv.org/abs/1807.11590
9 |
10 | ## Brief
11 |
12 | In short, Precise RoI Pooling is an integration-based (bilinear interpolation) average pooling method for RoI Pooling. It avoids any quantization and has a continuous gradient on bounding box coordinates. It is:
13 |
14 | - different from the original RoI Pooling proposed in [Fast R-CNN](https://arxiv.org/abs/1504.08083). PrRoI Pooling uses average pooling instead of max pooling for each bin and has a continuous gradient on bounding box coordinates. That is, one can take the derivatives of some loss function w.r.t the coordinates of each RoI and optimize the RoI coordinates.
15 | - different from the RoI Align proposed in [Mask R-CNN](https://arxiv.org/abs/1703.06870). PrRoI Pooling uses a full integration-based average pooling instead of sampling a constant number of points. This makes the gradient w.r.t. the coordinates continuous.
16 |
17 | For a better illustration, we illustrate RoI Pooling, RoI Align and PrRoI Pooing in the following figure. More details including the gradient computation can be found in our paper.
18 |
19 |
20 |
21 | ## Implementation
22 |
23 | PrRoI Pooling was originally implemented by [Tete Xiao](http://tetexiao.com/) based on MegBrain, an (internal) deep learning framework built by Megvii Inc. It was later adapted into open-source deep learning frameworks. Currently, we only support PyTorch. Unfortunately, we don't have any specific plan for the adaptation into other frameworks such as TensorFlow, but any contributions (pull requests) will be more than welcome.
24 |
25 | ## Usage (PyTorch 1.0)
26 |
27 | In the directory `pytorch/`, we provide a PyTorch-based implementation of PrRoI Pooling. It requires PyTorch 1.0+ and only supports CUDA (CPU mode is not implemented).
28 | Since we use PyTorch JIT for cxx/cuda code compilation, to use the module in your code, simply do:
29 |
30 | ```
31 | from prroi_pool import PrRoIPool2D
32 |
33 | avg_pool = PrRoIPool2D(window_height, window_width, spatial_scale)
34 | roi_features = avg_pool(features, rois)
35 |
36 | # for those who want to use the "functional"
37 |
38 | from prroi_pool.functional import prroi_pool2d
39 | roi_features = prroi_pool2d(features, rois, window_height, window_width, spatial_scale)
40 | ```
41 |
42 |
43 | ## Usage (PyTorch 0.4)
44 |
45 | **!!! Please first checkout to the branch pytorch0.4.**
46 |
47 | In the directory `pytorch/`, we provide a PyTorch-based implementation of PrRoI Pooling. It requires PyTorch 0.4 and only supports CUDA (CPU mode is not implemented).
48 | To use the PrRoI Pooling module, first goto `pytorch/prroi_pool` and execute `./travis.sh` to compile the essential components (you may need `nvcc` for this step). To use the module in your code, simply do:
49 |
50 | ```
51 | from prroi_pool import PrRoIPool2D
52 |
53 | avg_pool = PrRoIPool2D(window_height, window_width, spatial_scale)
54 | roi_features = avg_pool(features, rois)
55 |
56 | # for those who want to use the "functional"
57 |
58 | from prroi_pool.functional import prroi_pool2d
59 | roi_features = prroi_pool2d(features, rois, window_height, window_width, spatial_scale)
60 | ```
61 |
62 | Here,
63 |
64 | - RoI is an `m * 5` float tensor of format `(batch_index, x0, y0, x1, y1)`, following the convention in the original Caffe implementation of RoI Pooling, although in some frameworks the batch indices are provided by an integer tensor.
65 | - `spatial_scale` is multiplied to the RoIs. For example, if your feature maps are down-sampled by a factor of 16 (w.r.t. the input image), you should use a spatial scale of `1/16`.
66 | - The coordinates for RoI follows the [L, R) convension. That is, `(0, 0, 4, 4)` denotes a box of size `4x4`.
67 |
--------------------------------------------------------------------------------
/netdissect/upsegmodel/prroi_pool/__init__.py:
--------------------------------------------------------------------------------
1 | #! /usr/bin/env python3
2 | # -*- coding: utf-8 -*-
3 | # File : __init__.py
4 | # Author : Jiayuan Mao, Tete Xiao
5 | # Email : maojiayuan@gmail.com, jasonhsiao97@gmail.com
6 | # Date : 07/13/2018
7 | #
8 | # This file is part of PreciseRoIPooling.
9 | # Distributed under terms of the MIT license.
10 | # Copyright (c) 2017 Megvii Technology Limited.
11 |
12 | from .prroi_pool import *
13 |
14 |
--------------------------------------------------------------------------------
/netdissect/upsegmodel/prroi_pool/build.py:
--------------------------------------------------------------------------------
1 | #! /usr/bin/env python3
2 | # -*- coding: utf-8 -*-
3 | # File : build.py
4 | # Author : Jiayuan Mao, Tete Xiao
5 | # Email : maojiayuan@gmail.com, jasonhsiao97@gmail.com
6 | # Date : 07/13/2018
7 | #
8 | # This file is part of PreciseRoIPooling.
9 | # Distributed under terms of the MIT license.
10 | # Copyright (c) 2017 Megvii Technology Limited.
11 |
12 | import os
13 | import torch
14 |
15 | from torch.utils.ffi import create_extension
16 |
17 | headers = []
18 | sources = []
19 | defines = []
20 | extra_objects = []
21 | with_cuda = False
22 |
23 | if torch.cuda.is_available():
24 | with_cuda = True
25 |
26 | headers+= ['src/prroi_pooling_gpu.h']
27 | sources += ['src/prroi_pooling_gpu.c']
28 | defines += [('WITH_CUDA', None)]
29 |
30 | this_file = os.path.dirname(os.path.realpath(__file__))
31 | extra_objects_cuda = ['src/prroi_pooling_gpu_impl.cu.o']
32 | extra_objects_cuda = [os.path.join(this_file, fname) for fname in extra_objects_cuda]
33 | extra_objects.extend(extra_objects_cuda)
34 | else:
35 | # TODO(Jiayuan Mao @ 07/13): remove this restriction after we support the cpu implementation.
36 | raise NotImplementedError('Precise RoI Pooling only supports GPU (cuda) implememtations.')
37 |
38 | ffi = create_extension(
39 | '_prroi_pooling',
40 | headers=headers,
41 | sources=sources,
42 | define_macros=defines,
43 | relative_to=__file__,
44 | with_cuda=with_cuda,
45 | extra_objects=extra_objects
46 | )
47 |
48 | if __name__ == '__main__':
49 | ffi.build()
50 |
51 |
--------------------------------------------------------------------------------
/netdissect/upsegmodel/prroi_pool/functional.py:
--------------------------------------------------------------------------------
1 | #! /usr/bin/env python3
2 | # -*- coding: utf-8 -*-
3 | # File : functional.py
4 | # Author : Jiayuan Mao, Tete Xiao
5 | # Email : maojiayuan@gmail.com, jasonhsiao97@gmail.com
6 | # Date : 07/13/2018
7 | #
8 | # This file is part of PreciseRoIPooling.
9 | # Distributed under terms of the MIT license.
10 | # Copyright (c) 2017 Megvii Technology Limited.
11 |
12 | import torch
13 | import torch.autograd as ag
14 |
15 | try:
16 | from os.path import join as pjoin, dirname
17 | from torch.utils.cpp_extension import load as load_extension
18 | root_dir = pjoin(dirname(__file__), 'src')
19 | _prroi_pooling = load_extension(
20 | '_prroi_pooling',
21 | [pjoin(root_dir, 'prroi_pooling_gpu.c'), pjoin(root_dir, 'prroi_pooling_gpu_impl.cu')],
22 | verbose=False
23 | )
24 | except ImportError:
25 | raise ImportError('Can not compile Precise RoI Pooling library.')
26 |
27 | __all__ = ['prroi_pool2d']
28 |
29 |
30 | class PrRoIPool2DFunction(ag.Function):
31 | @staticmethod
32 | def forward(ctx, features, rois, pooled_height, pooled_width, spatial_scale):
33 | assert 'FloatTensor' in features.type() and 'FloatTensor' in rois.type(), \
34 | 'Precise RoI Pooling only takes float input, got {} for features and {} for rois.'.format(features.type(), rois.type())
35 |
36 | pooled_height = int(pooled_height)
37 | pooled_width = int(pooled_width)
38 | spatial_scale = float(spatial_scale)
39 |
40 | features = features.contiguous()
41 | rois = rois.contiguous()
42 | params = (pooled_height, pooled_width, spatial_scale)
43 |
44 | if features.is_cuda:
45 | output = _prroi_pooling.prroi_pooling_forward_cuda(features, rois, *params)
46 | ctx.params = params
47 | # everything here is contiguous.
48 | ctx.save_for_backward(features, rois, output)
49 | else:
50 | raise NotImplementedError('Precise RoI Pooling only supports GPU (cuda) implememtations.')
51 |
52 | return output
53 |
54 | @staticmethod
55 | def backward(ctx, grad_output):
56 | features, rois, output = ctx.saved_tensors
57 | grad_input = grad_coor = None
58 |
59 | if features.requires_grad:
60 | grad_output = grad_output.contiguous()
61 | grad_input = _prroi_pooling.prroi_pooling_backward_cuda(features, rois, output, grad_output, *ctx.params)
62 | if rois.requires_grad:
63 | grad_output = grad_output.contiguous()
64 | grad_coor = _prroi_pooling.prroi_pooling_coor_backward_cuda(features, rois, output, grad_output, *ctx.params)
65 |
66 | return grad_input, grad_coor, None, None, None
67 |
68 |
69 | prroi_pool2d = PrRoIPool2DFunction.apply
70 |
71 |
--------------------------------------------------------------------------------
/netdissect/upsegmodel/prroi_pool/prroi_pool.py:
--------------------------------------------------------------------------------
1 | #! /usr/bin/env python3
2 | # -*- coding: utf-8 -*-
3 | # File : prroi_pool.py
4 | # Author : Jiayuan Mao, Tete Xiao
5 | # Email : maojiayuan@gmail.com, jasonhsiao97@gmail.com
6 | # Date : 07/13/2018
7 | #
8 | # This file is part of PreciseRoIPooling.
9 | # Distributed under terms of the MIT license.
10 | # Copyright (c) 2017 Megvii Technology Limited.
11 |
12 | import torch.nn as nn
13 |
14 | from .functional import prroi_pool2d
15 |
16 | __all__ = ['PrRoIPool2D']
17 |
18 |
19 | class PrRoIPool2D(nn.Module):
20 | def __init__(self, pooled_height, pooled_width, spatial_scale):
21 | super().__init__()
22 |
23 | self.pooled_height = int(pooled_height)
24 | self.pooled_width = int(pooled_width)
25 | self.spatial_scale = float(spatial_scale)
26 |
27 | def forward(self, features, rois):
28 | return prroi_pool2d(features, rois, self.pooled_height, self.pooled_width, self.spatial_scale)
29 |
--------------------------------------------------------------------------------
/netdissect/upsegmodel/prroi_pool/src/prroi_pooling_gpu.c:
--------------------------------------------------------------------------------
1 | /*
2 | * File : prroi_pooling_gpu.c
3 | * Author : Jiayuan Mao, Tete Xiao
4 | * Email : maojiayuan@gmail.com, jasonhsiao97@gmail.com
5 | * Date : 07/13/2018
6 | *
7 | * Distributed under terms of the MIT license.
8 | * Copyright (c) 2017 Megvii Technology Limited.
9 | */
10 |
11 | #include
12 | #include
13 |
14 | #include
15 | #include
16 |
17 | #include
18 |
19 | #include "prroi_pooling_gpu_impl.cuh"
20 |
21 |
22 | at::Tensor prroi_pooling_forward_cuda(const at::Tensor &features, const at::Tensor &rois, int pooled_height, int pooled_width, float spatial_scale) {
23 | int nr_rois = rois.size(0);
24 | int nr_channels = features.size(1);
25 | int height = features.size(2);
26 | int width = features.size(3);
27 | int top_count = nr_rois * nr_channels * pooled_height * pooled_width;
28 | auto output = at::zeros({nr_rois, nr_channels, pooled_height, pooled_width}, features.options());
29 |
30 | if (output.numel() == 0) {
31 | THCudaCheck(cudaGetLastError());
32 | return output;
33 | }
34 |
35 | cudaStream_t stream = at::cuda::getCurrentCUDAStream();
36 | PrRoIPoolingForwardGpu(
37 | stream, features.data(), rois.data(), output.data(),
38 | nr_channels, height, width, pooled_height, pooled_width, spatial_scale,
39 | top_count
40 | );
41 |
42 | THCudaCheck(cudaGetLastError());
43 | return output;
44 | }
45 |
46 | at::Tensor prroi_pooling_backward_cuda(
47 | const at::Tensor &features, const at::Tensor &rois, const at::Tensor &output, const at::Tensor &output_diff,
48 | int pooled_height, int pooled_width, float spatial_scale) {
49 |
50 | auto features_diff = at::zeros_like(features);
51 |
52 | int nr_rois = rois.size(0);
53 | int batch_size = features.size(0);
54 | int nr_channels = features.size(1);
55 | int height = features.size(2);
56 | int width = features.size(3);
57 | int top_count = nr_rois * nr_channels * pooled_height * pooled_width;
58 | int bottom_count = batch_size * nr_channels * height * width;
59 |
60 | if (output.numel() == 0) {
61 | THCudaCheck(cudaGetLastError());
62 | return features_diff;
63 | }
64 |
65 | cudaStream_t stream = at::cuda::getCurrentCUDAStream();
66 | PrRoIPoolingBackwardGpu(
67 | stream,
68 | features.data(), rois.data(), output.data(), output_diff.data(),
69 | features_diff.data(),
70 | nr_channels, height, width, pooled_height, pooled_width, spatial_scale,
71 | top_count, bottom_count
72 | );
73 |
74 | THCudaCheck(cudaGetLastError());
75 | return features_diff;
76 | }
77 |
78 | at::Tensor prroi_pooling_coor_backward_cuda(
79 | const at::Tensor &features, const at::Tensor &rois, const at::Tensor &output, const at::Tensor &output_diff,
80 | int pooled_height, int pooled_width, float spatial_scale) {
81 |
82 | auto coor_diff = at::zeros_like(rois);
83 |
84 | int nr_rois = rois.size(0);
85 | int nr_channels = features.size(1);
86 | int height = features.size(2);
87 | int width = features.size(3);
88 | int top_count = nr_rois * nr_channels * pooled_height * pooled_width;
89 | int bottom_count = nr_rois * 5;
90 |
91 | if (output.numel() == 0) {
92 | THCudaCheck(cudaGetLastError());
93 | return coor_diff;
94 | }
95 |
96 | cudaStream_t stream = at::cuda::getCurrentCUDAStream();
97 | PrRoIPoolingCoorBackwardGpu(
98 | stream,
99 | features.data(), rois.data(), output.data(), output_diff.data(),
100 | coor_diff.data(),
101 | nr_channels, height, width, pooled_height, pooled_width, spatial_scale,
102 | top_count, bottom_count
103 | );
104 |
105 | THCudaCheck(cudaGetLastError());
106 | return coor_diff;
107 | }
108 |
109 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
110 | m.def("prroi_pooling_forward_cuda", &prroi_pooling_forward_cuda, "PRRoIPooling_forward");
111 | m.def("prroi_pooling_backward_cuda", &prroi_pooling_backward_cuda, "PRRoIPooling_backward");
112 | m.def("prroi_pooling_coor_backward_cuda", &prroi_pooling_coor_backward_cuda, "PRRoIPooling_backward_coor");
113 | }
114 |
--------------------------------------------------------------------------------
/netdissect/upsegmodel/prroi_pool/src/prroi_pooling_gpu.h:
--------------------------------------------------------------------------------
1 | /*
2 | * File : prroi_pooling_gpu.h
3 | * Author : Jiayuan Mao, Tete Xiao
4 | * Email : maojiayuan@gmail.com, jasonhsiao97@gmail.com
5 | * Date : 07/13/2018
6 | *
7 | * Distributed under terms of the MIT license.
8 | * Copyright (c) 2017 Megvii Technology Limited.
9 | */
10 |
11 | int prroi_pooling_forward_cuda(THCudaTensor *features, THCudaTensor *rois, THCudaTensor *output, int pooled_height, int pooled_width, float spatial_scale);
12 |
13 | int prroi_pooling_backward_cuda(
14 | THCudaTensor *features, THCudaTensor *rois, THCudaTensor *output, THCudaTensor *output_diff, THCudaTensor *features_diff,
15 | int pooled_height, int pooled_width, float spatial_scale
16 | );
17 |
18 | int prroi_pooling_coor_backward_cuda(
19 | THCudaTensor *features, THCudaTensor *rois, THCudaTensor *output, THCudaTensor *output_diff, THCudaTensor *features_diff,
20 | int pooled_height, int pooled_width, float spatial_scal
21 | );
22 |
23 |
--------------------------------------------------------------------------------
/netdissect/upsegmodel/prroi_pool/src/prroi_pooling_gpu_impl.cuh:
--------------------------------------------------------------------------------
1 | /*
2 | * File : prroi_pooling_gpu_impl.cuh
3 | * Author : Tete Xiao, Jiayuan Mao
4 | * Email : jasonhsiao97@gmail.com
5 | *
6 | * Distributed under terms of the MIT license.
7 | * Copyright (c) 2017 Megvii Technology Limited.
8 | */
9 |
10 | #ifndef PRROI_POOLING_GPU_IMPL_CUH
11 | #define PRROI_POOLING_GPU_IMPL_CUH
12 |
13 | #ifdef __cplusplus
14 | extern "C" {
15 | #endif
16 |
17 | #define F_DEVPTR_IN const float *
18 | #define F_DEVPTR_OUT float *
19 |
20 | void PrRoIPoolingForwardGpu(
21 | cudaStream_t stream,
22 | F_DEVPTR_IN bottom_data,
23 | F_DEVPTR_IN bottom_rois,
24 | F_DEVPTR_OUT top_data,
25 | const int channels_, const int height_, const int width_,
26 | const int pooled_height_, const int pooled_width_,
27 | const float spatial_scale_,
28 | const int top_count);
29 |
30 | void PrRoIPoolingBackwardGpu(
31 | cudaStream_t stream,
32 | F_DEVPTR_IN bottom_data,
33 | F_DEVPTR_IN bottom_rois,
34 | F_DEVPTR_IN top_data,
35 | F_DEVPTR_IN top_diff,
36 | F_DEVPTR_OUT bottom_diff,
37 | const int channels_, const int height_, const int width_,
38 | const int pooled_height_, const int pooled_width_,
39 | const float spatial_scale_,
40 | const int top_count, const int bottom_count);
41 |
42 | void PrRoIPoolingCoorBackwardGpu(
43 | cudaStream_t stream,
44 | F_DEVPTR_IN bottom_data,
45 | F_DEVPTR_IN bottom_rois,
46 | F_DEVPTR_IN top_data,
47 | F_DEVPTR_IN top_diff,
48 | F_DEVPTR_OUT bottom_diff,
49 | const int channels_, const int height_, const int width_,
50 | const int pooled_height_, const int pooled_width_,
51 | const float spatial_scale_,
52 | const int top_count, const int bottom_count);
53 |
54 | #ifdef __cplusplus
55 | } /* !extern "C" */
56 | #endif
57 |
58 | #endif /* !PRROI_POOLING_GPU_IMPL_CUH */
59 |
60 |
--------------------------------------------------------------------------------
/netdissect/upsegmodel/prroi_pool/test_prroi_pooling2d.py:
--------------------------------------------------------------------------------
1 | # -*- coding: utf-8 -*-
2 | # File : test_prroi_pooling2d.py
3 | # Author : Jiayuan Mao
4 | # Email : maojiayuan@gmail.com
5 | # Date : 18/02/2018
6 | #
7 | # This file is part of Jacinle.
8 |
9 | import unittest
10 |
11 | import torch
12 | import torch.nn as nn
13 | import torch.nn.functional as F
14 |
15 | from jactorch.utils.unittest import TorchTestCase
16 |
17 | from prroi_pool import PrRoIPool2D
18 |
19 |
20 | class TestPrRoIPool2D(TorchTestCase):
21 | def test_forward(self):
22 | pool = PrRoIPool2D(7, 7, spatial_scale=0.5)
23 | features = torch.rand((4, 16, 24, 32)).cuda()
24 | rois = torch.tensor([
25 | [0, 0, 0, 14, 14],
26 | [1, 14, 14, 28, 28],
27 | ]).float().cuda()
28 |
29 | out = pool(features, rois)
30 | out_gold = F.avg_pool2d(features, kernel_size=2, stride=1)
31 |
32 | self.assertTensorClose(out, torch.stack((
33 | out_gold[0, :, :7, :7],
34 | out_gold[1, :, 7:14, 7:14],
35 | ), dim=0))
36 |
37 | def test_backward_shapeonly(self):
38 | pool = PrRoIPool2D(2, 2, spatial_scale=0.5)
39 |
40 | features = torch.rand((4, 2, 24, 32)).cuda()
41 | rois = torch.tensor([
42 | [0, 0, 0, 4, 4],
43 | [1, 14, 14, 18, 18],
44 | ]).float().cuda()
45 | features.requires_grad = rois.requires_grad = True
46 | out = pool(features, rois)
47 |
48 | loss = out.sum()
49 | loss.backward()
50 |
51 | self.assertTupleEqual(features.size(), features.grad.size())
52 | self.assertTupleEqual(rois.size(), rois.grad.size())
53 |
54 |
55 | if __name__ == '__main__':
56 | unittest.main()
57 |
--------------------------------------------------------------------------------
/netdissect/upsegmodel/resnet.py:
--------------------------------------------------------------------------------
1 | import os
2 | import sys
3 | import torch
4 | import torch.nn as nn
5 | import math
6 | try:
7 | from lib.nn import SynchronizedBatchNorm2d
8 | except ImportError:
9 | from torch.nn import BatchNorm2d as SynchronizedBatchNorm2d
10 |
11 | try:
12 | from urllib import urlretrieve
13 | except ImportError:
14 | from urllib.request import urlretrieve
15 |
16 |
17 | __all__ = ['ResNet', 'resnet50', 'resnet101'] # resnet101 is coming soon!
18 |
19 |
20 | model_urls = {
21 | 'resnet50': 'http://sceneparsing.csail.mit.edu/model/pretrained_resnet/resnet50-imagenet.pth',
22 | 'resnet101': 'http://sceneparsing.csail.mit.edu/model/pretrained_resnet/resnet101-imagenet.pth'
23 | }
24 |
25 |
26 | def conv3x3(in_planes, out_planes, stride=1):
27 | "3x3 convolution with padding"
28 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
29 | padding=1, bias=False)
30 |
31 |
32 | class BasicBlock(nn.Module):
33 | expansion = 1
34 |
35 | def __init__(self, inplanes, planes, stride=1, downsample=None):
36 | super(BasicBlock, self).__init__()
37 | self.conv1 = conv3x3(inplanes, planes, stride)
38 | self.bn1 = SynchronizedBatchNorm2d(planes)
39 | self.relu = nn.ReLU(inplace=True)
40 | self.conv2 = conv3x3(planes, planes)
41 | self.bn2 = SynchronizedBatchNorm2d(planes)
42 | self.downsample = downsample
43 | self.stride = stride
44 |
45 | def forward(self, x):
46 | residual = x
47 |
48 | out = self.conv1(x)
49 | out = self.bn1(out)
50 | out = self.relu(out)
51 |
52 | out = self.conv2(out)
53 | out = self.bn2(out)
54 |
55 | if self.downsample is not None:
56 | residual = self.downsample(x)
57 |
58 | out += residual
59 | out = self.relu(out)
60 |
61 | return out
62 |
63 |
64 | class Bottleneck(nn.Module):
65 | expansion = 4
66 |
67 | def __init__(self, inplanes, planes, stride=1, downsample=None):
68 | super(Bottleneck, self).__init__()
69 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
70 | self.bn1 = SynchronizedBatchNorm2d(planes)
71 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
72 | padding=1, bias=False)
73 | self.bn2 = SynchronizedBatchNorm2d(planes)
74 | self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
75 | self.bn3 = SynchronizedBatchNorm2d(planes * 4)
76 | self.relu = nn.ReLU(inplace=True)
77 | self.downsample = downsample
78 | self.stride = stride
79 |
80 | def forward(self, x):
81 | residual = x
82 |
83 | out = self.conv1(x)
84 | out = self.bn1(out)
85 | out = self.relu(out)
86 |
87 | out = self.conv2(out)
88 | out = self.bn2(out)
89 | out = self.relu(out)
90 |
91 | out = self.conv3(out)
92 | out = self.bn3(out)
93 |
94 | if self.downsample is not None:
95 | residual = self.downsample(x)
96 |
97 | out += residual
98 | out = self.relu(out)
99 |
100 | return out
101 |
102 |
103 | class ResNet(nn.Module):
104 |
105 | def __init__(self, block, layers, num_classes=1000):
106 | self.inplanes = 128
107 | super(ResNet, self).__init__()
108 | self.conv1 = conv3x3(3, 64, stride=2)
109 | self.bn1 = SynchronizedBatchNorm2d(64)
110 | self.relu1 = nn.ReLU(inplace=True)
111 | self.conv2 = conv3x3(64, 64)
112 | self.bn2 = SynchronizedBatchNorm2d(64)
113 | self.relu2 = nn.ReLU(inplace=True)
114 | self.conv3 = conv3x3(64, 128)
115 | self.bn3 = SynchronizedBatchNorm2d(128)
116 | self.relu3 = nn.ReLU(inplace=True)
117 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
118 |
119 | self.layer1 = self._make_layer(block, 64, layers[0])
120 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
121 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
122 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
123 | self.avgpool = nn.AvgPool2d(7, stride=1)
124 | self.fc = nn.Linear(512 * block.expansion, num_classes)
125 |
126 | for m in self.modules():
127 | if isinstance(m, nn.Conv2d):
128 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
129 | m.weight.data.normal_(0, math.sqrt(2. / n))
130 | elif isinstance(m, SynchronizedBatchNorm2d):
131 | m.weight.data.fill_(1)
132 | m.bias.data.zero_()
133 |
134 | def _make_layer(self, block, planes, blocks, stride=1):
135 | downsample = None
136 | if stride != 1 or self.inplanes != planes * block.expansion:
137 | downsample = nn.Sequential(
138 | nn.Conv2d(self.inplanes, planes * block.expansion,
139 | kernel_size=1, stride=stride, bias=False),
140 | SynchronizedBatchNorm2d(planes * block.expansion),
141 | )
142 |
143 | layers = []
144 | layers.append(block(self.inplanes, planes, stride, downsample))
145 | self.inplanes = planes * block.expansion
146 | for i in range(1, blocks):
147 | layers.append(block(self.inplanes, planes))
148 |
149 | return nn.Sequential(*layers)
150 |
151 | def forward(self, x):
152 | x = self.relu1(self.bn1(self.conv1(x)))
153 | x = self.relu2(self.bn2(self.conv2(x)))
154 | x = self.relu3(self.bn3(self.conv3(x)))
155 | x = self.maxpool(x)
156 |
157 | x = self.layer1(x)
158 | x = self.layer2(x)
159 | x = self.layer3(x)
160 | x = self.layer4(x)
161 |
162 | x = self.avgpool(x)
163 | x = x.view(x.size(0), -1)
164 | x = self.fc(x)
165 |
166 | return x
167 |
168 | '''
169 | def resnet18(pretrained=False, **kwargs):
170 | """Constructs a ResNet-18 model.
171 |
172 | Args:
173 | pretrained (bool): If True, returns a model pre-trained on Places
174 | """
175 | model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
176 | if pretrained:
177 | model.load_state_dict(load_url(model_urls['resnet18']))
178 | return model
179 |
180 |
181 | def resnet34(pretrained=False, **kwargs):
182 | """Constructs a ResNet-34 model.
183 |
184 | Args:
185 | pretrained (bool): If True, returns a model pre-trained on Places
186 | """
187 | model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
188 | if pretrained:
189 | model.load_state_dict(load_url(model_urls['resnet34']))
190 | return model
191 | '''
192 |
193 | def resnet50(pretrained=False, **kwargs):
194 | """Constructs a ResNet-50 model.
195 |
196 | Args:
197 | pretrained (bool): If True, returns a model pre-trained on Places
198 | """
199 | model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
200 | if pretrained:
201 | model.load_state_dict(load_url(model_urls['resnet50']), strict=False)
202 | return model
203 |
204 |
205 | def resnet101(pretrained=False, **kwargs):
206 | """Constructs a ResNet-101 model.
207 |
208 | Args:
209 | pretrained (bool): If True, returns a model pre-trained on Places
210 | """
211 | model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
212 | if pretrained:
213 | model.load_state_dict(load_url(model_urls['resnet101']), strict=False)
214 | return model
215 |
216 | # def resnet152(pretrained=False, **kwargs):
217 | # """Constructs a ResNet-152 model.
218 | #
219 | # Args:
220 | # pretrained (bool): If True, returns a model pre-trained on Places
221 | # """
222 | # model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
223 | # if pretrained:
224 | # model.load_state_dict(load_url(model_urls['resnet152']))
225 | # return model
226 |
227 | def load_url(url, model_dir='./pretrained', map_location=None):
228 | if not os.path.exists(model_dir):
229 | os.makedirs(model_dir)
230 | filename = url.split('/')[-1]
231 | cached_file = os.path.join(model_dir, filename)
232 | if not os.path.exists(cached_file):
233 | sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file))
234 | urlretrieve(url, cached_file)
235 | return torch.load(cached_file, map_location=map_location)
236 |
--------------------------------------------------------------------------------
/netdissect/upsegmodel/resnext.py:
--------------------------------------------------------------------------------
1 | import os
2 | import sys
3 | import torch
4 | import torch.nn as nn
5 | import math
6 | try:
7 | from lib.nn import SynchronizedBatchNorm2d
8 | except ImportError:
9 | from torch.nn import BatchNorm2d as SynchronizedBatchNorm2d
10 |
11 | try:
12 | from urllib import urlretrieve
13 | except ImportError:
14 | from urllib.request import urlretrieve
15 |
16 |
17 | __all__ = ['ResNeXt', 'resnext101'] # support resnext 101
18 |
19 |
20 | model_urls = {
21 | #'resnext50': 'http://sceneparsing.csail.mit.edu/model/pretrained_resnet/resnext50-imagenet.pth',
22 | 'resnext101': 'http://sceneparsing.csail.mit.edu/model/pretrained_resnet/resnext101-imagenet.pth'
23 | }
24 |
25 |
26 | def conv3x3(in_planes, out_planes, stride=1):
27 | "3x3 convolution with padding"
28 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
29 | padding=1, bias=False)
30 |
31 |
32 | class GroupBottleneck(nn.Module):
33 | expansion = 2
34 |
35 | def __init__(self, inplanes, planes, stride=1, groups=1, downsample=None):
36 | super(GroupBottleneck, self).__init__()
37 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
38 | self.bn1 = SynchronizedBatchNorm2d(planes)
39 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
40 | padding=1, groups=groups, bias=False)
41 | self.bn2 = SynchronizedBatchNorm2d(planes)
42 | self.conv3 = nn.Conv2d(planes, planes * 2, kernel_size=1, bias=False)
43 | self.bn3 = SynchronizedBatchNorm2d(planes * 2)
44 | self.relu = nn.ReLU(inplace=True)
45 | self.downsample = downsample
46 | self.stride = stride
47 |
48 | def forward(self, x):
49 | residual = x
50 |
51 | out = self.conv1(x)
52 | out = self.bn1(out)
53 | out = self.relu(out)
54 |
55 | out = self.conv2(out)
56 | out = self.bn2(out)
57 | out = self.relu(out)
58 |
59 | out = self.conv3(out)
60 | out = self.bn3(out)
61 |
62 | if self.downsample is not None:
63 | residual = self.downsample(x)
64 |
65 | out += residual
66 | out = self.relu(out)
67 |
68 | return out
69 |
70 |
71 | class ResNeXt(nn.Module):
72 |
73 | def __init__(self, block, layers, groups=32, num_classes=1000):
74 | self.inplanes = 128
75 | super(ResNeXt, self).__init__()
76 | self.conv1 = conv3x3(3, 64, stride=2)
77 | self.bn1 = SynchronizedBatchNorm2d(64)
78 | self.relu1 = nn.ReLU(inplace=True)
79 | self.conv2 = conv3x3(64, 64)
80 | self.bn2 = SynchronizedBatchNorm2d(64)
81 | self.relu2 = nn.ReLU(inplace=True)
82 | self.conv3 = conv3x3(64, 128)
83 | self.bn3 = SynchronizedBatchNorm2d(128)
84 | self.relu3 = nn.ReLU(inplace=True)
85 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
86 |
87 | self.layer1 = self._make_layer(block, 128, layers[0], groups=groups)
88 | self.layer2 = self._make_layer(block, 256, layers[1], stride=2, groups=groups)
89 | self.layer3 = self._make_layer(block, 512, layers[2], stride=2, groups=groups)
90 | self.layer4 = self._make_layer(block, 1024, layers[3], stride=2, groups=groups)
91 | self.avgpool = nn.AvgPool2d(7, stride=1)
92 | self.fc = nn.Linear(1024 * block.expansion, num_classes)
93 |
94 | for m in self.modules():
95 | if isinstance(m, nn.Conv2d):
96 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels // m.groups
97 | m.weight.data.normal_(0, math.sqrt(2. / n))
98 | elif isinstance(m, SynchronizedBatchNorm2d):
99 | m.weight.data.fill_(1)
100 | m.bias.data.zero_()
101 |
102 | def _make_layer(self, block, planes, blocks, stride=1, groups=1):
103 | downsample = None
104 | if stride != 1 or self.inplanes != planes * block.expansion:
105 | downsample = nn.Sequential(
106 | nn.Conv2d(self.inplanes, planes * block.expansion,
107 | kernel_size=1, stride=stride, bias=False),
108 | SynchronizedBatchNorm2d(planes * block.expansion),
109 | )
110 |
111 | layers = []
112 | layers.append(block(self.inplanes, planes, stride, groups, downsample))
113 | self.inplanes = planes * block.expansion
114 | for i in range(1, blocks):
115 | layers.append(block(self.inplanes, planes, groups=groups))
116 |
117 | return nn.Sequential(*layers)
118 |
119 | def forward(self, x):
120 | x = self.relu1(self.bn1(self.conv1(x)))
121 | x = self.relu2(self.bn2(self.conv2(x)))
122 | x = self.relu3(self.bn3(self.conv3(x)))
123 | x = self.maxpool(x)
124 |
125 | x = self.layer1(x)
126 | x = self.layer2(x)
127 | x = self.layer3(x)
128 | x = self.layer4(x)
129 |
130 | x = self.avgpool(x)
131 | x = x.view(x.size(0), -1)
132 | x = self.fc(x)
133 |
134 | return x
135 |
136 |
137 | '''
138 | def resnext50(pretrained=False, **kwargs):
139 | """Constructs a ResNet-50 model.
140 |
141 | Args:
142 | pretrained (bool): If True, returns a model pre-trained on Places
143 | """
144 | model = ResNeXt(GroupBottleneck, [3, 4, 6, 3], **kwargs)
145 | if pretrained:
146 | model.load_state_dict(load_url(model_urls['resnext50']), strict=False)
147 | return model
148 | '''
149 |
150 |
151 | def resnext101(pretrained=False, **kwargs):
152 | """Constructs a ResNet-101 model.
153 |
154 | Args:
155 | pretrained (bool): If True, returns a model pre-trained on Places
156 | """
157 | model = ResNeXt(GroupBottleneck, [3, 4, 23, 3], **kwargs)
158 | if pretrained:
159 | model.load_state_dict(load_url(model_urls['resnext101']), strict=False)
160 | return model
161 |
162 |
163 | # def resnext152(pretrained=False, **kwargs):
164 | # """Constructs a ResNeXt-152 model.
165 | #
166 | # Args:
167 | # pretrained (bool): If True, returns a model pre-trained on Places
168 | # """
169 | # model = ResNeXt(GroupBottleneck, [3, 8, 36, 3], **kwargs)
170 | # if pretrained:
171 | # model.load_state_dict(load_url(model_urls['resnext152']))
172 | # return model
173 |
174 |
175 | def load_url(url, model_dir='./pretrained', map_location=None):
176 | if not os.path.exists(model_dir):
177 | os.makedirs(model_dir)
178 | filename = url.split('/')[-1]
179 | cached_file = os.path.join(model_dir, filename)
180 | if not os.path.exists(cached_file):
181 | sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file))
182 | urlretrieve(url, cached_file)
183 | return torch.load(cached_file, map_location=map_location)
184 |
--------------------------------------------------------------------------------
/netdissect/workerpool.py:
--------------------------------------------------------------------------------
1 | '''
2 | WorkerPool and WorkerBase for handling the common problems in managing
3 | a multiprocess pool of workers that aren't done by multiprocessing.Pool,
4 | including setup with per-process state, debugging by putting the worker
5 | on the main thread, and correct handling of unexpected errors, and ctrl-C.
6 |
7 | To use it,
8 | 1. Put the per-process setup and the per-task work in the
9 | setup() and work() methods of your own WorkerBase subclass.
10 | 2. To prepare the process pool, instantiate a WorkerPool, passing your
11 | subclass type as the first (worker) argument, as well as any setup keyword
12 | arguments. The WorkerPool will instantiate one of your workers in each
13 | worker process (passing in the setup arguments in those processes).
14 | If debugging, the pool can have process_count=0 to force all the work
15 | to be done immediately on the main thread; otherwise all the work
16 | will be passed to other processes.
17 | 3. Whenever there is a new piece of work to distribute, call pool.add(*args).
18 | The arguments will be queued and passed as worker.work(*args) to the
19 | next available worker.
20 | 4. When all the work has been distributed, call pool.join() to wait for all
21 | the work to complete and to finish and terminate all the worker processes.
22 | When pool.join() returns, all the work will have been done.
23 |
24 | No arrangement is made to collect the results of the work: for example,
25 | the return value of work() is ignored. If you need to collect the
26 | results, use your own mechanism (filesystem, shared memory object, queue)
27 | which can be distributed using setup arguments.
28 | '''
29 |
30 | from multiprocessing import Process, Queue, cpu_count
31 | import signal
32 | import atexit
33 | import sys
34 |
35 |
36 | class WorkerBase(Process):
37 | '''
38 | Subclass this class and override its work() method (and optionally,
39 | setup() as well) to define the units of work to be done in a process
40 | worker in a woker pool.
41 | '''
42 |
43 | def __init__(self, i, process_count, queue, initargs):
44 | if process_count > 0:
45 | # Make sure we ignore ctrl-C if we are not on main process.
46 | signal.signal(signal.SIGINT, signal.SIG_IGN)
47 | self.process_id = i
48 | self.process_count = process_count
49 | self.queue = queue
50 | super(WorkerBase, self).__init__()
51 | self.setup(**initargs)
52 |
53 | def run(self):
54 | # Do the work until None is dequeued
55 | while True:
56 | try:
57 | work_batch = self.queue.get()
58 | except (KeyboardInterrupt, SystemExit):
59 | print('Exiting...')
60 | break
61 | if work_batch is None:
62 | self.queue.put(None) # for another worker
63 | return
64 | self.work(*work_batch)
65 |
66 | def setup(self, **initargs):
67 | '''
68 | Override this method for any per-process initialization.
69 | Keywoard args are passed from WorkerPool constructor.
70 | '''
71 | pass
72 |
73 | def work(self, *args):
74 | '''
75 | Override this method for one-time initialization.
76 | Args are passed from WorkerPool.add() arguments.
77 | '''
78 | raise NotImplementedError('worker subclass needed')
79 |
80 |
81 | class WorkerPool(object):
82 | '''
83 | Instantiate this object (passing a WorkerBase subclass type
84 | as its first argument) to create a worker pool. Then call
85 | pool.add(*args) to queue args to distribute to worker.work(*args),
86 | and call pool.join() to wait for all the workers to complete.
87 | '''
88 |
89 | def __init__(self, worker=WorkerBase, process_count=None, **initargs):
90 | global active_pools
91 | if process_count is None:
92 | process_count = cpu_count()
93 | if process_count == 0:
94 | # zero process_count uses only main process, for debugging.
95 | self.queue = None
96 | self.processes = None
97 | self.worker = worker(None, 0, None, initargs)
98 | return
99 | # Ctrl-C strategy: worker processes should ignore ctrl-C. Set
100 | # this up to be inherited by child processes before forking.
101 | original_sigint_handler = signal.signal(signal.SIGINT, signal.SIG_IGN)
102 | active_pools[id(self)] = self
103 | self.queue = Queue(maxsize=(process_count * 3))
104 | self.processes = None # Initialize before trying to construct workers
105 | self.processes = [worker(i, process_count, self.queue, initargs)
106 | for i in range(process_count)]
107 | for p in self.processes:
108 | p.start()
109 | # The main process should handle ctrl-C. Restore this now.
110 | signal.signal(signal.SIGINT, original_sigint_handler)
111 |
112 | def add(self, *work_batch):
113 | if self.queue is None:
114 | if hasattr(self, 'worker'):
115 | self.worker.work(*work_batch)
116 | else:
117 | print('WorkerPool shutting down.', file=sys.stderr)
118 | else:
119 | try:
120 | # The queue can block if the work is so slow it gets full.
121 | self.queue.put(work_batch)
122 | except (KeyboardInterrupt, SystemExit):
123 | # Handle ctrl-C if done while waiting for the queue.
124 | self.early_terminate()
125 |
126 | def join(self):
127 | # End the queue, and wait for all worker processes to complete nicely.
128 | if self.queue is not None:
129 | self.queue.put(None)
130 | for p in self.processes:
131 | p.join()
132 | self.queue = None
133 | # Remove myself from the set of pools that need cleanup on shutdown.
134 | try:
135 | del active_pools[id(self)]
136 | except:
137 | pass
138 |
139 | def early_terminate(self):
140 | # When shutting down unexpectedly, first end the queue.
141 | if self.queue is not None:
142 | try:
143 | self.queue.put_nowait(None) # Nonblocking put throws if full.
144 | self.queue = None
145 | except:
146 | pass
147 | # But then don't wait: just forcibly terminate workers.
148 | if self.processes is not None:
149 | for p in self.processes:
150 | p.terminate()
151 | self.processes = None
152 | try:
153 | del active_pools[id(self)]
154 | except:
155 | pass
156 |
157 | def __del__(self):
158 | if self.queue is not None:
159 | print('ERROR: workerpool.join() not called!', file=sys.stderr)
160 | self.join()
161 |
162 |
163 | # Error and ctrl-C handling: kill worker processes if the main process ends.
164 | active_pools = {}
165 |
166 |
167 | def early_terminate_pools():
168 | for _, pool in list(active_pools.items()):
169 | pool.early_terminate()
170 |
171 |
172 | atexit.register(early_terminate_pools)
173 |
--------------------------------------------------------------------------------
/netdissect/zdataset.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import numpy
3 | import itertools
4 | from torch.utils.data import TensorDataset
5 |
6 |
7 | def z_dataset_for_model(model, size=100, seed=1, indices=None):
8 | if indices is not None:
9 | indices = torch.as_tensor(indices, dtype=torch.int64, device='cpu')
10 | zs = z_sample_for_model(model, indices.max().item() + 1, seed)
11 | zs = zs[indices]
12 | else:
13 | zs = z_sample_for_model(model, size, seed)
14 | return TensorDataset(zs)
15 |
16 |
17 | def z_sample_for_model(model, size=100, seed=1):
18 | # If the model is marked with an input shape, use it.
19 | if hasattr(model, 'input_shape'):
20 | sample = standard_z_sample(size, model.input_shape[1], seed=seed).view(
21 | (size,) + model.input_shape[1:])
22 | return sample
23 | # Examine first conv in model to determine input feature size.
24 | first_layer = [c for c in model.modules()
25 | if isinstance(c, (torch.nn.Conv2d, torch.nn.ConvTranspose2d,
26 | torch.nn.Linear))][0]
27 | # 4d input if convolutional, 2d input if first layer is linear.
28 | if isinstance(first_layer, (torch.nn.Conv2d, torch.nn.ConvTranspose2d)):
29 | sample = standard_z_sample(
30 | size, first_layer.in_channels, seed=seed)[:, :, None, None]
31 | else:
32 | sample = standard_z_sample(
33 | size, first_layer.in_features, seed=seed)
34 | return sample
35 |
36 |
37 | def standard_z_sample(size, depth, seed=1, device=None):
38 | '''
39 | Generate a standard set of random Z as a (size, z_dimension) tensor.
40 | With the same random seed, it always returns the same z (e.g.,
41 | the first one is always the same regardless of the size.)
42 | '''
43 | # Use numpy RandomState since it can be done deterministically
44 | # without affecting global state
45 | rng = numpy.random.RandomState(seed)
46 | result = torch.from_numpy(
47 | rng.standard_normal(size * depth)
48 | .reshape(size, depth)).float()
49 | if device is not None:
50 | result = result.to(device)
51 | return result
52 |
53 |
54 | def standard_y_sample(size, num_classes, seed=1, device=None):
55 | '''
56 | Generate a standard set of random categorical as a (size,) tensor
57 | of integers up to (num_classes-1).
58 | With the same random seed, it always returns the same y (e.g.,
59 | the first one is always the same regardless of the size.)
60 | '''
61 | # Use numpy RandomState since it can be done deterministically
62 | # without affecting global state
63 | rng = numpy.random.RandomState(seed)
64 | result = torch.from_numpy(
65 | rng.randint(num_classes, size=size)).long()
66 | if device is not None:
67 | result = result.to(device)
68 | return result
69 |
70 |
71 | def training_loader(z_generator, batch_size, loader_size=10000):
72 | '''
73 | Returns an infinite generator that runs through randomized z
74 | batches, forever.
75 | '''
76 | g_epoch = 1
77 | while True:
78 | z_data = z_dataset_for_model(
79 | z_generator, size=loader_size, seed=g_epoch + 1)
80 | dataloader = torch.utils.data.DataLoader(
81 | z_data,
82 | shuffle=False,
83 | batch_size=batch_size,
84 | num_workers=10,
85 | pin_memory=True)
86 | for batch in dataloader:
87 | yield batch
88 | g_epoch += 1
89 |
90 |
91 | def testing_loader(z_generator, batch_size, test_size=1000):
92 | '''
93 | Returns an a short iterator that returns a small set of test data.
94 | '''
95 | z_data = z_dataset_for_model(
96 | z_generator, size=test_size, seed=1)
97 | dataloader = torch.utils.data.DataLoader(
98 | z_data,
99 | shuffle=False,
100 | batch_size=batch_size,
101 | num_workers=10,
102 | pin_memory=True)
103 | return dataloader
104 |
105 |
106 | def epoch_grouper(loader, epoch_size):
107 | '''
108 | To use with the infinite training loader: groups the training data
109 | batches into epochs of the given size.
110 | '''
111 | it = iter(loader)
112 | while True:
113 | chunk_it = itertools.islice(it, epoch_size)
114 | try:
115 | first_el = next(chunk_it)
116 | except StopIteration:
117 | return
118 | yield itertools.chain((first_el,), chunk_it)
119 |
--------------------------------------------------------------------------------
/notebooks/ipynb_drop_output.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 |
3 | """
4 | Suppress output and prompt numbers in git version control.
5 |
6 | This script will tell git to ignore prompt numbers and cell output
7 | when looking at ipynb files UNLESS their metadata contains:
8 |
9 | "git" : { "keep_output" : true }
10 |
11 | The notebooks themselves are not changed.
12 |
13 | See also this blogpost: http://pascalbugnion.net/blog/ipython-notebooks-and-git.html.
14 |
15 | Usage instructions
16 | ==================
17 |
18 | 1. Put this script in a directory that is on the system's path.
19 | For future reference, I will assume you saved it in
20 | `~/scripts/ipynb_drop_output`.
21 | 2. Make sure it is executable by typing the command
22 | `chmod +x ~/scripts/ipynb_drop_output`.
23 | 3. Register a filter for ipython notebooks by
24 | putting the following line in `~/.config/git/attributes`:
25 | `*.ipynb filter=clean_ipynb`
26 | 4. Connect this script to the filter by running the following
27 | git commands:
28 |
29 | git config --global filter.clean_ipynb.clean ipynb_drop_output
30 | git config --global filter.clean_ipynb.smudge cat
31 |
32 | To tell git NOT to ignore the output and prompts for a notebook,
33 | open the notebook's metadata (Edit > Edit Notebook Metadata). A
34 | panel should open containing the lines:
35 |
36 | {
37 | "name" : "",
38 | "signature" : "some very long hash"
39 | }
40 |
41 | Add an extra line so that the metadata now looks like:
42 |
43 | {
44 | "name" : "",
45 | "signature" : "don't change the hash, but add a comma at the end of the line",
46 | "git" : { "keep_outputs" : true }
47 | }
48 |
49 | You may need to "touch" the notebooks for git to actually register a change, if
50 | your notebooks are already under version control.
51 |
52 | Notes
53 | =====
54 |
55 | Changed by David Bau to make stripping output the default.
56 |
57 | This script is inspired by http://stackoverflow.com/a/20844506/827862, but
58 | lets the user specify whether the ouptut of a notebook should be kept
59 | in the notebook's metadata, and works for IPython v3.0.
60 | """
61 |
62 | import sys
63 | import json
64 |
65 | nb = sys.stdin.read()
66 |
67 | json_in = json.loads(nb)
68 | nb_metadata = json_in["metadata"]
69 | keep_output = False
70 | if "git" in nb_metadata:
71 | if "keep_outputs" in nb_metadata["git"] and nb_metadata["git"]["keep_outputs"]:
72 | keep_output = True
73 | if keep_output:
74 | sys.stdout.write(nb)
75 | exit()
76 |
77 |
78 | ipy_version = int(json_in["nbformat"])-1 # nbformat is 1 more than actual version.
79 |
80 | def strip_output_from_cell(cell):
81 | if "outputs" in cell:
82 | cell["outputs"] = []
83 | if "prompt_number" in cell:
84 | del cell["prompt_number"]
85 | if "execution_count" in cell:
86 | cell["execution_count"] = None
87 |
88 |
89 | if ipy_version == 2:
90 | for sheet in json_in["worksheets"]:
91 | for cell in sheet["cells"]:
92 | strip_output_from_cell(cell)
93 | else:
94 | for cell in json_in["cells"]:
95 | strip_output_from_cell(cell)
96 |
97 | json.dump(json_in, sys.stdout, sort_keys=True, indent=1, separators=(",",": "))
98 |
--------------------------------------------------------------------------------
/notebooks/setup_notebooks.sh:
--------------------------------------------------------------------------------
1 | #!/bin/bash
2 |
3 | # Start from directory of script
4 | cd "$(dirname "${BASH_SOURCE[0]}")"
5 |
6 | # Set up git config filters so huge output of notebooks is not committed.
7 | git config filter.clean_ipynb.clean "$(pwd)/ipynb_drop_output.py"
8 | git config filter.clean_ipynb.smudge cat
9 | git config filter.clean_ipynb.required true
10 |
11 | # Set up symlinks for the example notebooks
12 | for DIRNAME in datasets results netdissect experiment
13 | do
14 | ln -sfn ../${DIRNAME} .
15 | done
16 |
--------------------------------------------------------------------------------
/notebooks/shapebias_experiment.ipynb:
--------------------------------------------------------------------------------
1 | {
2 | "cells": [
3 | {
4 | "cell_type": "code",
5 | "execution_count": null,
6 | "metadata": {},
7 | "outputs": [],
8 | "source": [
9 | "from netdissect import parallelfolder, show, tally, nethook, renormalize\n",
10 | "from experiment import readdissect, setting\n",
11 | "import copy, PIL.Image\n",
12 | "from netdissect import upsample, imgsave, imgviz\n",
13 | "import re, torchvision, torch, os\n",
14 | "from IPython.display import SVG\n",
15 | "from matplotlib import pyplot as plt\n",
16 | "\n",
17 | "def normalize_filename(n):\n",
18 | " return re.match(r'^(.*Places365_\\w+_\\d+)', n).group(1)\n",
19 | "\n",
20 | "ds = parallelfolder.ParallelImageFolders(\n",
21 | " ['datasets/places/val', 'datasets/stylized-places/val'],\n",
22 | " transform=torchvision.transforms.Compose([\n",
23 | " torchvision.transforms.Resize(256),\n",
24 | " # transforms.CenterCrop(224),\n",
25 | " torchvision.transforms.CenterCrop(256),\n",
26 | " torchvision.transforms.ToTensor(),\n",
27 | " renormalize.NORMALIZER['imagenet'],\n",
28 | " ]),\n",
29 | " normalize_filename=normalize_filename,\n",
30 | " shuffle=True)\n",
31 | "\n",
32 | "\n",
33 | "layers = [\n",
34 | " 'conv5_3',\n",
35 | " 'conv5_2',\n",
36 | " 'conv5_1',\n",
37 | " 'conv4_3',\n",
38 | " 'conv4_2',\n",
39 | " 'conv4_1',\n",
40 | " 'conv3_3',\n",
41 | " 'conv3_2',\n",
42 | " 'conv3_1',\n",
43 | " 'conv2_2',\n",
44 | " 'conv2_1',\n",
45 | " 'conv1_2',\n",
46 | " 'conv1_1',\n",
47 | "]\n",
48 | "qd = readdissect.DissectVis(layers=layers)\n",
49 | "net = setting.load_classifier('vgg16')\n",
50 | "\n",
51 | "sds = parallelfolder.ParallelImageFolders(\n",
52 | " ['datasets/stylized-places/val'],\n",
53 | " transform=torchvision.transforms.Compose([\n",
54 | " torchvision.transforms.Resize(256),\n",
55 | " # transforms.CenterCrop(224),\n",
56 | " torchvision.transforms.CenterCrop(256),\n",
57 | " torchvision.transforms.ToTensor(),\n",
58 | " renormalize.NORMALIZER['imagenet'],\n",
59 | " ]),\n",
60 | " normalize_filename=normalize_filename,\n",
61 | " shuffle=True)\n",
62 | "\n",
63 | "def s_image(layername, unit):\n",
64 | " result = PIL.Image.open(os.path.join(qd.dir(layername), 's_imgs/unit%d.jpg' % unit))\n",
65 | " result.load()\n",
66 | " return result\n",
67 | "\n",
68 | "\n",
69 | "def su_image(layername, unit):\n",
70 | " result = PIL.Image.open(os.path.join(qd.dir(layername), 'su_imgs/unit%d.jpg' % unit))\n",
71 | " result.load()\n",
72 | " return result\n"
73 | ]
74 | },
75 | {
76 | "cell_type": "code",
77 | "execution_count": null,
78 | "metadata": {
79 | "scrolled": false
80 | },
81 | "outputs": [],
82 | "source": [
83 | "for layername in layers:\n",
84 | " inst_net = nethook.InstrumentedModel(copy.deepcopy(net)).cuda()\n",
85 | " inst_net.retain_layer('features.' + layername)\n",
86 | " inst_net(ds[0][0][None].cuda())\n",
87 | " sample_act = inst_net.retained_layer('features.' + layername).cpu()\n",
88 | " upfn = upsample.upsampler((64, 64), sample_act.shape[2:])\n",
89 | "\n",
90 | " def flat_acts(batch):\n",
91 | " inst_net(batch.cuda())\n",
92 | " acts = upfn(inst_net.retained_layer('features.' + layername))\n",
93 | " return acts.permute(0, 2, 3, 1).contiguous().view(-1, acts.shape[1])\n",
94 | " s_rq = tally.tally_quantile(flat_acts, sds, cachefile=os.path.join(qd.dir(layername), 's_rq.npz'))\n",
95 | " u_rq = qd.rq(layername)\n",
96 | "\n",
97 | " def intersect_99_fn(uimg, simg):\n",
98 | " s_99 = s_rq.quantiles(0.99)[None,:,None,None].cuda()\n",
99 | " u_99 = u_rq.quantiles(0.99)[None,:,None,None].cuda()\n",
100 | " with torch.no_grad():\n",
101 | " ux, sx = uimg.cuda(), simg.cuda()\n",
102 | " inst_net(ux)\n",
103 | " ur = inst_net.retained_layer('features.' + layername)\n",
104 | " inst_net(sx)\n",
105 | " sr = inst_net.retained_layer('features.' + layername)\n",
106 | " return ((sr > s_99).float() * (ur > u_99).float()).permute(0, 2, 3, 1).reshape(-1, ur.size(1))\n",
107 | " \n",
108 | " intersect_99 = tally.tally_mean(intersect_99_fn, ds,\n",
109 | " cachefile=os.path.join(qd.dir(layername), 'intersect_99.npz'))\n",
110 | " print(layername)\n",
111 | " numerator = intersect_99.mean()\n",
112 | " denominator = (0.02 - intersect_99.mean())\n",
113 | " score = (numerator / denominator).clamp(0, 1)\n",
114 | " plt.plot(score)\n",
115 | " plt.show()\n",
116 | " fig, ax = plt.subplots(1, 1, figsize=(3,1.2), dpi=300)\n",
117 | " ax.hist(score)\n",
118 | " ax.set_ylabel('%s units' % (layername.replace('features.', '')))\n",
119 | " ax.spines['right'].set_visible(False)\n",
120 | " ax.spines['top'].set_visible(False)\n",
121 | " # ax.set_xlabel('unit IoU (stylized vs original)')\n",
122 | " plt.show()\n",
123 | " labelcat_list_h = []\n",
124 | " labelcat_list_l = []\n",
125 | " for i, rec in enumerate(qd.labels[layername]):\n",
126 | " if rec['iou'] and float(rec['iou']) >= 0.04:\n",
127 | " if score[i] > 0.1:\n",
128 | " labelcat_list_h.append((rec['label'], rec['cat']))\n",
129 | " else:\n",
130 | " labelcat_list_l.append((rec['label'], rec['cat']))\n",
131 | " display(SVG(qd.bargraph_from_conceptcatlist(labelcat_list_l)))\n",
132 | " display(SVG(qd.bargraph_from_conceptcatlist(labelcat_list_h)))\n",
133 | " \n",
134 | " ordering = score.sort()[1]\n",
135 | "\n",
136 | " for i in torch.cat([ordering[:5], ordering[-10:]]):\n",
137 | " #if qd.iou(layername, i) > 0.04:\n",
138 | " print(i.item(), score[i].item(), qd.label(layername, i), qd.iou(layername, i))\n",
139 | " display(qd.image(layername, i))\n",
140 | " display(s_image(layername, i))\n",
141 | "\n",
142 | " #result = [qd.iou(layername, i) for i in ordering]\n",
143 | " #plt.plot(result)"
144 | ]
145 | },
146 | {
147 | "cell_type": "code",
148 | "execution_count": null,
149 | "metadata": {
150 | "scrolled": false
151 | },
152 | "outputs": [],
153 | "source": [
154 | "fig, axes = plt.subplots(5, 1, figsize=(5,6), dpi=300, sharex=True)\n",
155 | "plotlayers = [\n",
156 | " 'features.conv1_2',\n",
157 | " 'features.conv2_2',\n",
158 | " 'features.conv3_3',\n",
159 | " 'features.conv4_3',\n",
160 | " 'features.conv5_3',\n",
161 | "]\n",
162 | "for i, layername in enumerate(plotlayers):\n",
163 | " inst_net = nethook.InstrumentedModel(copy.deepcopy(net)).cuda()\n",
164 | " inst_net.retain_layer(layername)\n",
165 | " \n",
166 | " def flat_acts(batch):\n",
167 | " inst_net(batch.cuda())\n",
168 | " acts = upfn(inst_net.retained_layer(layername))\n",
169 | " return acts.permute(0, 2, 3, 1).contiguous().view(-1, acts.shape[1])\n",
170 | " s_rq = tally.tally_quantile(flat_acts, sds, cachefile=os.path.join(qd.dir(layername), 's_rq.npz'))\n",
171 | " u_rq = qd.rq(layername)\n",
172 | "\n",
173 | " def intersect_99_fn(uimg, simg):\n",
174 | " s_99 = s_rq.quantiles(0.99)[None,:,None,None].cuda()\n",
175 | " u_99 = u_rq.quantiles(0.99)[None,:,None,None].cuda()\n",
176 | " with torch.no_grad():\n",
177 | " ux, sx = uimg.cuda(), simg.cuda()\n",
178 | " inst_net(ux)\n",
179 | " ur = inst_net.retained_layer(layername)\n",
180 | " inst_net(sx)\n",
181 | " sr = inst_net.retained_layer(layername)\n",
182 | " return ((sr > s_99).float() * (ur > u_99).float()).permute(0, 2, 3, 1).reshape(-1, ur.size(1))\n",
183 | " \n",
184 | " intersect_99 = tally.tally_mean(intersect_99_fn, ds,\n",
185 | " cachefile=os.path.join(qd.dir(layername), 'intersect_99.npz'))\n",
186 | " numerator = intersect_99.mean()\n",
187 | " denominator = (0.02 - intersect_99.mean())\n",
188 | " score = (numerator / denominator).clamp(0, 0.5)\n",
189 | " ax = axes[i]\n",
190 | " ax.hist(score)\n",
191 | " # ax.set_ylabel('%s' % (layername.replace('features.', '')))\n",
192 | " ax.spines['right'].set_visible(False)\n",
193 | " ax.spines['top'].set_visible(False)\n",
194 | " # ax.set_xlabel('unit IoU (stylized vs original)')\n",
195 | "plt.show()\n"
196 | ]
197 | },
198 | {
199 | "cell_type": "code",
200 | "execution_count": null,
201 | "metadata": {},
202 | "outputs": [],
203 | "source": [
204 | "for u in [166, 107, 268, 434, 436, 437, 73, 220, 299, 494, 485, 477, 462, 338]:\n",
205 | " print(u, score[u].item())"
206 | ]
207 | },
208 | {
209 | "cell_type": "code",
210 | "execution_count": null,
211 | "metadata": {},
212 | "outputs": [],
213 | "source": [
214 | "qd.dirs"
215 | ]
216 | }
217 | ],
218 | "metadata": {
219 | "kernelspec": {
220 | "display_name": "Python 3",
221 | "language": "python",
222 | "name": "python3"
223 | },
224 | "language_info": {
225 | "codemirror_mode": {
226 | "name": "ipython",
227 | "version": 3
228 | },
229 | "file_extension": ".py",
230 | "mimetype": "text/x-python",
231 | "name": "python",
232 | "nbconvert_exporter": "python",
233 | "pygments_lexer": "ipython3",
234 | "version": "3.6.10"
235 | }
236 | },
237 | "nbformat": 4,
238 | "nbformat_minor": 4
239 | }
--------------------------------------------------------------------------------
/setup/denv.yml:
--------------------------------------------------------------------------------
1 | name: denv
2 | channels:
3 | - pytorch
4 | - ostrokach-forge
5 | - conda-forge
6 | dependencies:
7 | - python=3.6
8 | - cudatoolkit=10.2
9 | - cudnn=7.6.0
10 | - pytorch=1.5.1
11 | - torchvision
12 | - mkl-include
13 | - numpy
14 | - scipy
15 | - scikit-learn
16 | - scikit-image
17 | - matplotlib
18 | - seaborn
19 | - numba
20 | - jupyter
21 | - jupyterlab
22 | - pyyaml
23 | - mkl
24 | - tqdm
25 | - pip
26 |
--------------------------------------------------------------------------------
/setup/setup_denv.sh:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env bash
2 |
3 | # Bash script to set up an anaconda python-based deep learning environment
4 | # that has support for pytorch, tensorflow, pycaffe in the same environment,
5 | # long with juypter, scipy etc.
6 |
7 | # This should not require root. However, it does copy and build a lot of
8 | # binaries into your ~/.conda directory. If you do not want to store
9 | # these in your homedir disk, then ~/.conda can be a symlink somewhere else.
10 | # (At MIT CSAIL, you should symlink ~/.conda to a directory on NFS or local
11 | # disk instead of leaving it on AFS, or else you will exhaust your quota.)
12 |
13 | # Start from parent directory of script
14 | cd "$(dirname "$(dirname "$(readlink -f "$0")")")"
15 |
16 | # Default RECIPE 'denv' can be overridden by 'RECIPE=foo setup.sh'
17 | RECIPE=${RECIPE:-denv}
18 | # Default ENV_NAME 'denv' can be overridden by 'ENV_NAME=foo setup.sh'
19 | ENV_NAME="${ENV_NAME:-${RECIPE}}"
20 | echo "Creating conda environment ${ENV_NAME}"
21 |
22 | if [[ ! $(type -P conda) ]]
23 | then
24 | echo "conda not in PATH"
25 | echo "read: https://conda.io/docs/user-guide/install/index.html"
26 | exit 1
27 | fi
28 |
29 | if df "${HOME}/.conda" --type=afs > /dev/null 2>&1
30 | then
31 | echo "Not installing: your ~/.conda directory is on AFS."
32 | echo "Use 'ln -s /some/nfs/dir ~/.conda' to avoid using up your AFS quota."
33 | exit 1
34 | fi
35 |
36 | # Uninstall existing environment
37 | source deactivate
38 | rm -rf ~/.conda/envs/${ENV_NAME}
39 |
40 | # Build new environment: torch and torch vision from source
41 | # CUDA_HOME is needed
42 | # https://github.com/rusty1s/pytorch_scatter/issues/19#issuecomment-449735614
43 | conda env create --name=${ENV_NAME} -f setup/${RECIPE}.yml
44 |
45 | # Set up CUDA_HOME to set itself up correctly on every source activate
46 | # https://stackoverflow.com/questions/31598963
47 | mkdir -p ~/.conda/envs/${ENV_NAME}/etc/conda/activate.d
48 | echo "export CUDA_HOME=/usr/local/cuda-10.2" > \
49 | ~/.conda/envs/${ENV_NAME}/etc/conda/activate.d/CUDA_HOME.sh
50 |
51 | source activate ${ENV_NAME}
52 |
53 |
--------------------------------------------------------------------------------
/stylization/LICENSE:
--------------------------------------------------------------------------------
1 | MIT License
2 |
3 | Copyright (c) 2020 MIT CSAIL and David Bau
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.
22 |
--------------------------------------------------------------------------------
/stylization/README.md:
--------------------------------------------------------------------------------
1 | # stylize-datasets
2 | This repository contains code for stylizing arbitrary image datasets using [AdaIN](https://arxiv.org/abs/1703.06868). The code is a generalization of Robert Geirhos' [Stylized-ImageNet](https://github.com/rgeirhos/Stylized-ImageNet) code, which is tailored to stylizing ImageNet. Everything in this repository is based on naoto0804's [pytorch-AdaIN](https://github.com/naoto0804/pytorch-AdaIN) implementation.
3 |
4 | Given an image dataset, the script creates the specified number of stylized versions of every image while keeping the directory structure and naming scheme intact (usefull for existing data loaders or if directory names include class annotations).
5 |
6 | Feel free to open an issue in case there is any question.
7 |
8 | ## Usage
9 | - Dependencies:
10 | - python >= 3.6
11 | - Pillow
12 | - torch
13 | - torchvision
14 | - tqdm
15 | - Download the models:
16 | - either run run `bash models/download_models.sh` or download the models manually from [vgg](https://drive.google.com/file/d/108uza-dsmwvbW2zv-G73jtVcMU_2Nb7Y/view)/[decoder](https://drive.google.com/file/d/1w9r1NoYnn7tql1VYG3qDUzkbIks24RBQ/view) and move both files to the `models/` directory
17 | - Get style images: Download train.zip from [Kaggle's painter-by-numbers dataset](https://www.kaggle.com/c/painter-by-numbers/data)
18 | - To stylize a dataset, run `python stylize.py`.
19 |
20 | Arguments:
21 | - `--content-dir ` the top-level directory of the content image dataset (mandatory)
22 | - `--style-dir ` the top-level directory of the style images (mandatory)
23 | - `--output-dir
