Recycling3

robustness/attacker.py

@@ -271,6 +271,7 @@ class AttackerModel(ch.nn.Module):
     def __init__(self, model, dataset):
         super(AttackerModel, self).__init__()
         self.normalizer = helpers.InputNormalize(dataset.mean, dataset.std)
+        self.normalizer = ch.jit.script(self.normalizer)
         self.model = model
 
     def forward(self, inp):

robustness/imagenet_models/resnet.py

@@ -7,7 +7,7 @@
 __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
            'resnet152', 'resnext50_32x4d', 'resnext101_32x8d',
            'wide_resnet50_2', 'wide_resnet101_2',
-           'wide_resnet50_3', 'wide_resnet50_4', 'wide_resnet50_5', 
+           'wide_resnet50_3', 'wide_resnet50_4', 'wide_resnet50_5',
            'wide_resnet50_6', ]
 
 
@@ -57,7 +57,7 @@ def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
         self.downsample = downsample
         self.stride = stride
 
-    def forward(self, x, fake_relu=False, no_relu=False):
+    def forward(self, x):
         identity = x
 
         out = self.conv1(x)
@@ -72,10 +72,6 @@ def forward(self, x, fake_relu=False, no_relu=False):
 
         out += identity
 
-        if fake_relu:
-            return FakeReLU.apply(out)
-        if no_relu:
-            return out
         return self.relu(out)
 
 
@@ -201,10 +197,10 @@ def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
                                 base_width=self.base_width, dilation=self.dilation,
                                 norm_layer=norm_layer))
 
-        return SequentialWithArgs(*layers)
+        return nn.Sequential(*layers)
         # return nn.Sequential(*layers)
 
-    def _forward(self, x, with_latent=False, fake_relu=False, no_relu=False):
+    def _forward(self, x):
         x = self.conv1(x)
         x = self.bn1(x)
         x = self.relu(x)
@@ -213,13 +209,11 @@ def _forward(self, x, with_latent=False, fake_relu=False, no_relu=False):
         x = self.layer1(x)
         x = self.layer2(x)
         x = self.layer3(x)
-        x = self.layer4(x, fake_relu=fake_relu, no_relu=no_relu)
+        x = self.layer4(x)
 
         x = self.avgpool(x)
         pre_out = torch.flatten(x, 1)
         final = self.fc(pre_out)
-        if with_latent:
-            return final, pre_out
         return final
 
     # Allow for accessing forward method in a inherited class
@@ -360,7 +354,7 @@ def wide_resnet50_3(pretrained=False, progress=True, **kwargs):
 
 
 def wide_resnet50_4(pretrained=False, progress=True, **kwargs):
-    r"""Wide ResNet-50-4 model 
+    r"""Wide ResNet-50-4 model
     Args:
         pretrained (bool): If True, returns a model pre-trained on ImageNet
         progress (bool): If True, displays a progress bar of the download to stderr

robustness/main.py

@@ -56,7 +56,7 @@ def main(args, store=None):
 
     if not args.resume_optimizer: checkpoint = None
     model = train_model(args, model, loaders, store=store,
-                                    checkpoint=checkpoint)
+                        checkpoint=checkpoint, val_loader=val_loader)
     return model
 
 def setup_args(args):

robustness/tools/custom_modules.py

@@ -1,4 +1,4 @@
-import torch 
+import torch
 from torch import nn
 ch = torch
 
@@ -16,12 +16,12 @@ def forward(self, x):
         return FakeReLU.apply(x)
 
 class SequentialWithArgs(torch.nn.Sequential):
-    def forward(self, input, *args, **kwargs):
+    def forward(self, input):
         vs = list(self._modules.values())
         l = len(vs)
         for i in range(l):
             if i == l-1:
-                input = vs[i](input, *args, **kwargs)
+                input = vs[i](input)
             else:
                 input = vs[i](input)
         return input

robustness/tools/helpers.py

@@ -96,63 +96,141 @@ def __init__(self, new_mean, new_std):
         self.register_buffer("new_std", new_std)
 
     def forward(self, x):
-        x = ch.clamp(x, 0, 1)
+        # x = ch.clamp(x, 0, 1)
         x_normalized = (x - self.new_mean)/self.new_std
         return x_normalized
 
+# class DataPrefetcher():
+#     def __init__(self, loader, stop_after=None):
+#         self.loader = loader
+#         self.dataset = loader.dataset
+
+#     def __len__(self):
+#         return len(self.loader)
+
+#     def __iter__(self):
+#         # count = 0
+#         # self.loaditer = iter(self.loader)
+#         # self.preload()
+#         # for i in range(loaditer)
+#         # while self.next_input is not None:
+#         #     ch.cuda.current_stream().wait_stream(self.stream)
+#         #     input = self.next_input
+#         #     target = self.next_target
+#         #     self.preload()
+#         #     count += 1
+#         #     yield input, target
+#         #     if type(self.stop_after) is int and (count > self.stop_after):
+#         #         break
+#         prev_x, prev_y = None, None
+#         for _, (x, y) in enumerate(self.loader):
+#             x = x.to(device='cuda', memory_format=ch.channels_last,
+#                      non_blocking=True).to(dtype=ch.float32, non_blocking=True)
+#             y = y.to(device='cuda', non_blocking=True)
+#             if prev_x is None:
+#                 prev_x, prev_y = x, y
+#             else:
+#                 yield prev_x, prev_y
+#                 prev_x, prev_y = x, y
+
+#         yield prev_x, prev_y
+
 class DataPrefetcher():
     def __init__(self, loader, stop_after=None):
         self.loader = loader
-        self.dataset = loader.dataset
-        self.stream = ch.cuda.Stream()
-        self.stop_after = stop_after
-        self.next_input = None
-        self.next_target = None
+        self.dataset = self.loader.dataset
 
     def __len__(self):
         return len(self.loader)
 
+    def __iter__(self):
+        prefetcher = data_prefetcher(self.loader)
+        x, y = prefetcher.next()
+        while x is not None:
+            yield x, y
+            x, y = prefetcher.next()
+
+class data_prefetcher():
+    def __init__(self, loader):
+        self.loader = iter(loader)
+        self.stream = ch.cuda.Stream()
+        # With Amp, it isn't necessary to manually convert data to half.
+        # if args.fp16:
+        #     self.mean = self.mean.half()
+        #     self.std = self.std.half()
+        self.preload()
+
     def preload(self):
         try:
-            self.next_input, self.next_target = next(self.loaditer)
+            self.next_input, self.next_target = next(self.loader)
         except StopIteration:
             self.next_input = None
             self.next_target = None
             return
+        # if record_stream() doesn't work, another option is to make sure device inputs are created
+        # on the main stream.
+        # self.next_input_gpu = ch.empty_like(self.next_input, device='cuda')
+        # self.next_target_gpu = ch.empty_like(self.next_target, device='cuda')
+        # Need to make sure the memory allocated for next_* is not still in use by the main stream
+        # at the time we start copying to next_*:
+        # self.stream.wait_stream(ch.cuda.current_stream())
         with ch.cuda.stream(self.stream):
-            self.next_input = self.next_input.cuda(non_blocking=True)
+            # self.next_input = self.next_input.cuda(non_blocking=True)
             self.next_target = self.next_target.cuda(non_blocking=True)
+            # more code for the alternative if record_stream() doesn't work:
+            # copy_ will record the use of the pinned source tensor in this side stream.
+            # self.next_input_gpu.copy_(self.next_input, non_blocking=True)
+            # self.next_target_gpu.copy_(self.next_target, non_blocking=True)
+            # self.next_input = self.next_input_gpu
+            # self.next_target = self.next_target_gpu
 
-    def __iter__(self):
-        count = 0
-        self.loaditer = iter(self.loader)
+            # With Amp, it isn't necessary to manually convert data to half.
+            # if args.fp16:
+            #     self.next_input = self.next_input.half()
+            # else:
+            self.next_input = self.next_input.to(device='cuda', memory_format=ch.channels_last,
+                                                 non_blocking=True).to(dtype=ch.float32, non_blocking=True)
+            # to(dtype=ch.float32, non_blocking=True)
+
+    def next(self):
+        ch.cuda.current_stream().wait_stream(self.stream)
+        input = self.next_input
+        target = self.next_target
+        if input is not None:
+            input.record_stream(ch.cuda.current_stream())
+        if target is not None:
+            target.record_stream(ch.cuda.current_stream())
         self.preload()
-        while self.next_input is not None:
-            ch.cuda.current_stream().wait_stream(self.stream)
-            input = self.next_input
-            target = self.next_target
-            self.preload()
-            count += 1
-            yield input, target
-            if type(self.stop_after) is int and (count > self.stop_after):
-                break
+        return input, target
+
 
 class AverageMeter(object):
     """Computes and stores the average and current value"""
-    def __init__(self):
+    def __init__(self, name, fmt=':f'):
+        self.name = name
+        self.fmt = fmt
         self.reset()
 
     def reset(self):
         self.val = 0
-        self.avg = 0
         self.sum = 0
         self.count = 0
 
-    def update(self, val, n=1):
+    @property
+    def avg(self):
+        return self.sum / max(self.count, 1)
+
+    def update(self, val, _=0):
         self.val = val
-        self.sum += val * n
-        self.count += n
-        self.avg = self.sum / self.count
+        self.sum += val
+        self.count += 1
+        # self.avg = self.sum / self.count
+
+    def __str__(self):
+        fmtstr = '{name} {val' + self.fmt + '} ({avg2' + self.fmt + '} ({sum' + self.fmt + '})'
+        self.avg2 = self.avg
+        # self.avg = (self.sum / max(self.count, 1))
+        return fmtstr.format(**self.__dict__)
 
 # ImageNet label mappings
 def get_label_mapping(dataset_name, ranges):