├── .gitignore ├── LICENSE ├── README.md ├── RayS.py ├── RayS_Single.py ├── arch ├── cifar_model.py ├── fs_utils.py ├── madry_wrn.py ├── mnist_model.py ├── preact_resnet.py ├── wideresnet.py ├── wideresnet_compact.py ├── wideresnet_fs.py ├── wideresnet_he.py ├── wideresnet_interp.py ├── wideresnet_overfitting.py └── wideresnet_rst.py ├── attack_natural.py ├── attack_robust.py ├── dataset.py ├── general_tf_model.py ├── general_torch_model.py ├── images └── adbd_leaderboard.png ├── model ├── cifar10_gpu.pt └── mnist_gpu.pt └── pgbar.py /.gitignore: -------------------------------------------------------------------------------- 1 | *.txt 2 | *.sh 3 | model/rob* 4 | model/*.dict 5 | data/ 6 | __pycache__/ 7 | .DS_Store 8 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 Jinghui Chen, Quanquan Gu 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 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # RayS: A Ray Searching Method for Hard-label Adversarial Attack (KDD2020) 2 | "RayS: A Ray Searching Method for Hard-label Adversarial Attack"\ 3 | *[Jinghui Chen](https://jinghuichen.github.io)*, *[Quanquan Gu](http://web.cs.ucla.edu/~qgu/)*\ 4 | [https://arxiv.org/abs/2006.12792](https://arxiv.org/abs/2006.12792) 5 | 6 | This repository contains our PyTorch implementation of RayS: A Ray Searching Method for Hard-label Adversarial Attack in the paper [RayS: A Ray Searching Method for Hard-label Adversarial Attack](https://arxiv.org/abs/2006.12792) (accepted by KDD 2020). 7 | 8 | # What is RayS 9 | RayS is a hard-label adversarial attack which only requires the target model's hard-label output (prediction label). 10 | 11 | **It is gradient-free, hyper-parameter free, and is also independent of adversarial losses such as CrossEntropy or C&W.** 12 | 13 | Therefore, RayS can be used as a good sanity check for possible "falsely robust" models (models that may overfit to certain types of gradient-based attacks and adversarial losses). 14 | 15 | # Average Decision Boundary Distance (ADBD) 16 | RayS also proposed a new model robustness metric: `ADBD` (average decision boundary distance), which reflects examples' average distance to their closest decision boundary. 17 | 18 | # Model Robustness: ADBD Leaderboard 19 | 20 | We tested the robustness of recently proposed robust models which are trained on the CIFAR-10 dataset with the maximum L_inf norm perturbation strength `epsilon=0.031` (8/255). The robustness is evaluated on the entire CIFAR-10 testset (10000 examples). 21 | 22 | **Note**: 23 | * Ranking is based on the ADBD (average decision boundary distance) metric under RayS attack with the default query limit set as 40000. Reducing the query limit will accelerate the process but may lead to inaccurate ADBD value. For fast checking purpose, we recommend evaluating on subset of CIFAR-10 testset (e.g., 1000 examples). 24 | * `*` denotes model using extra data for training. 25 | * `Robust Acc (RayS)` represents robust accuracy under RayS attack for L_inf norm perturbation strength `epsilon=0.031` (8/255). For truly robust models, this value could be larger than the reported value (using white-box attacks) due to the hard-label limitation. For the current best robust accuracy evaluation, please refers to [AutoAttack](https://github.com/fra31/auto-attack), which uses an ensemble of four white-box/black-box attacks. 26 | * `ADBD` represents our proposed Average Decision Boundary Distance metric, which is independent to the perturbation strength `epsilon`. It reflects the overall model robustness through the lens of decision boundary distance. `ADBD` can be served as a complement to the traditional robust accuracy metric. Furthermore, `ADBD` only depends on hard-label output and can be adopted for cases where back-propgation or even soft-labels are not available. 27 | 28 | 29 | |Method |Natural Acc |Robust Acc
(Reported) |Robust Acc
(RayS) | ADBD| 30 | |:---:|---:|---:|---:|---:| 31 | | [WAR
(Wu et al., 2020)](https://arxiv.org/abs/2010.01279)*| 85.6| 59.8| 63.2| 0.0480| 32 | | [RST
(Carmon et al., 2019)](https://arxiv.org/abs/1905.13736)*| 89.7| 62.5| 64.6| 0.0465| 33 | | [HYDRA
(Sehwag et al., 2020)](https://arxiv.org/abs/2002.10509)*| 89.0| 57.2| 62.1| 0.0450| 34 | | [MART
(Wang et al., 2020)](https://openreview.net/forum?id=rklOg6EFwS)*| 87.5| 65.0| 62.2| 0.0439| 35 | | [UAT++
(Alayrac et al., 2019)](https://arxiv.org/abs/1905.13725)*| 86.5| 56.3| 62.1| 0.0426| 36 | | [Pretraining
(Hendrycks et al., 2019)](https://arxiv.org/abs/1901.09960)*| 87.1| 57.4| 60.1| 0.0419| 37 | | [Robust-overfitting
(Rice et al., 2020)](https://arxiv.org/abs/2002.11569)| 85.3 | 58.0 | 58.6 | 0.0404| 38 | | [TRADES
(Zhang et al., 2019b)](https://arxiv.org/abs/1901.08573)| 85.4| 56.4| 57.3| 0.0403| 39 | | [Backward Smoothing
(Chen et al., 2020)](https://arxiv.org/abs/2010.01278)| 85.3| 54.9| 55.1| 0.0403| 40 | | [Adversarial Training (retrained)
(Madry et al., 2018)](https://arxiv.org/abs/1706.06083)| 87.4| 50.6| 54.0| 0.0377| 41 | | [MMA
(Ding et al., 2020)](https://openreview.net/forum?id=HkeryxBtPB)| 84.4| 47.2| 47.7| 0.0345| 42 | | [Adversarial Training (original)
(Madry et al., 2018)](https://arxiv.org/abs/1706.06083)| 87.1| 47.0| 50.7| 0.0344| 43 | | [Fast Adversarial Training
(Wong et al., 2020)](https://arxiv.org/abs/2001.03994)| 83.8| 46.1| 50.1 | 0.0334| 44 | | [Adv-Interp
(Zhang & Xu, 2020)](https://openreview.net/forum?id=Syejj0NYvr¬eId=Syejj0NYvr) | 91.0| 68.7| 46.9| 0.0305| 45 | | [Feature-Scatter
(Zhang & Wang, 2019)](http://papers.nips.cc/paper/8459-defense-against-adversarial-attacks-using-feature-scattering-based-adversarial-training)| 91.3| 60.6| 44.5| 0.0301| 46 | | [SENSE
(Kim & Wang, 2020)](https://openreview.net/forum?id=rJlf_RVKwr)| 91.9| 57.2| 43.9| 0.0288| 47 | 48 | 49 |

50 | 51 |

52 | 53 | 54 | Please contact us if you want to add your model to the leaderboard. 55 | 56 | 57 | ## How to use RayS to evaluate your model robustness: 58 | 59 | ### Prerequisites: 60 | * Python 61 | * Numpy 62 | * CUDA 63 | 64 | ### PyTorch models 65 | Import RayS attack by 66 | 67 | ```python 68 | from general_torch_model import GeneralTorchModel 69 | torch_model = GeneralTorchModel(model, n_class=10, im_mean=None, im_std=None) 70 | 71 | from RayS import RayS 72 | attack = RayS(torch_model, epsilon=args.epsilon) 73 | ``` 74 | 75 | where: 76 | + `torch_model` is the PyTorch model under GeneralTorchModel warpper; For models using transformed images (exceed the range of [0,1]), simply set `im_mean=[0.5, 0.5, 0.5]` and `im_std=[0.5, 0.5, 0.5]` for instance, 77 | + `epsilon` is the maximum adversarial perturbation strength. 78 | 79 | To actually run RayS attack, use 80 | 81 | ```python 82 | x_adv, queries, adbd, succ = attack(data, label, query_limit) 83 | ``` 84 | 85 | it returns: 86 | + `x_adv`: the adversarial examples found by RayS, 87 | + `queries`: the number of queries used for finding the adversarial examples, 88 | + `adbd`: the average decision boundary distance for each example, 89 | + `succ`: indicate whether each example being successfully attacked. 90 | 91 | 92 | * Sample usage on attacking a robust model: 93 | ```bash 94 | - python3 attack_robust.py --dataset rob_cifar_trades --query 40000 --batch 1000 --epsilon 0.031 95 | ``` 96 | + You can also use `--num 1000` argument to limit the number of examples to be attacked as 1000. Default `num` is set as 10000 (the whole CIFAR10 testset). 97 | 98 | ### TensorFlow models 99 | To evaluate TensorFlow models with RayS attack: 100 | 101 | ```python 102 | from general_tf_model import GeneralTFModel 103 | tf_model = GeneralTFModel(model.logits, model.x_input, sess, n_class=10, im_mean=None, im_std=None) 104 | 105 | from RayS import RayS 106 | attack = RayS(tf_model, epsilon=args.epsilon) 107 | ``` 108 | 109 | where: 110 | + `model.logits`: logits tensor return by the Tensorflow model, 111 | + `model.x_input`: placeholder for model input (NHWC format), 112 | + `sess`: TF session . 113 | 114 | The remaining part is the same as evaluating PyTorch models. 115 | 116 | 117 | 118 | ## Reproduce experiments in the paper: 119 | * Run attacks on a naturally trained model (Inception): 120 | ```bash 121 | - python3 attack_natural.py --dataset inception --epsilon 0.05 122 | ``` 123 | * Run attacks on a naturally trained model (Resnet): 124 | ```bash 125 | - python3 attack_natural.py --dataset resnet --epsilon 0.05 126 | ``` 127 | * Run attacks on a naturally trained model (Cifar): 128 | ```bash 129 | - python3 attack_natural.py --dataset cifar --epsilon 0.031 130 | ``` 131 | * Run attacks on a naturally trained model (MNIST): 132 | ```bash 133 | - python3 attack_natural.py --dataset mnist --epsilon 0.3 134 | ``` 135 | 136 | 137 | ## Citation 138 | Please check our paper for technical details and full results. 139 | 140 | ``` 141 | @inproceedings{chen2020rays, 142 | title={RayS: A Ray Searching Method for Hard-label Adversarial Attack}, 143 | author={Chen, Jinghui and Gu, Quanquan}, 144 | booktitle={Proceedings of the 26rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining}, 145 | year={2020} 146 | } 147 | ``` 148 | 149 | ## Contact 150 | If you have any question regarding RayS attack or the ADBD leaderboard above, please contact jinghuic@ucla.edu, enjoy! -------------------------------------------------------------------------------- /RayS.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | from pgbar import progress_bar 4 | 5 | 6 | class RayS(object): 7 | def __init__(self, model, epsilon=0.031, order=np.inf): 8 | self.model = model 9 | self.ord = order 10 | self.epsilon = epsilon 11 | self.sgn_t = None 12 | self.d_t = None 13 | self.x_final = None 14 | self.queries = None 15 | 16 | def get_xadv(self, x, v, d, lb=0., ub=1.): 17 | if isinstance(d, int): 18 | d = torch.tensor(d).repeat(len(x)).cuda() 19 | out = x + d.view(len(x), 1, 1, 1) * v 20 | out = torch.clamp(out, lb, ub) 21 | return out 22 | 23 | def attack_hard_label(self, x, y, target=None, query_limit=10000, seed=None): 24 | """ Attack the original image and return adversarial example 25 | model: (pytorch model) 26 | (x, y): original image 27 | """ 28 | shape = list(x.shape) 29 | dim = np.prod(shape[1:]) 30 | if seed is not None: 31 | np.random.seed(seed) 32 | 33 | # init variables 34 | self.queries = torch.zeros_like(y).cuda() 35 | self.sgn_t = torch.sign(torch.ones(shape)).cuda() 36 | self.d_t = torch.ones_like(y).float().fill_(float("Inf")).cuda() 37 | working_ind = (self.d_t > self.epsilon).nonzero().flatten() 38 | 39 | stop_queries = self.queries.clone() 40 | dist = self.d_t.clone() 41 | self.x_final = self.get_xadv(x, self.sgn_t, self.d_t) 42 | 43 | block_level = 0 44 | block_ind = 0 45 | for i in range(query_limit): 46 | block_num = 2 ** block_level 47 | block_size = int(np.ceil(dim / block_num)) 48 | start, end = block_ind * block_size, min(dim, (block_ind + 1) * block_size) 49 | 50 | valid_mask = (self.queries < query_limit) 51 | attempt = self.sgn_t.clone().view(shape[0], dim) 52 | attempt[valid_mask.nonzero().flatten(), start:end] *= -1. 53 | attempt = attempt.view(shape) 54 | 55 | self.binary_search(x, y, target, attempt, valid_mask) 56 | 57 | block_ind += 1 58 | if block_ind == 2 ** block_level or end == dim: 59 | block_level += 1 60 | block_ind = 0 61 | 62 | dist = torch.norm((self.x_final - x).view(shape[0], -1), self.ord, 1) 63 | stop_queries[working_ind] = self.queries[working_ind] 64 | working_ind = (dist > self.epsilon).nonzero().flatten() 65 | 66 | if torch.sum(self.queries >= query_limit) == shape[0]: 67 | print('out of queries') 68 | break 69 | 70 | progress_bar(torch.min(self.queries.float()), query_limit, 71 | 'd_t: %.4f | adbd: %.4f | queries: %.4f | rob acc: %.4f | iter: %d' 72 | % (torch.mean(self.d_t), torch.mean(dist), torch.mean(self.queries.float()), 73 | len(working_ind) / len(x), i + 1)) 74 | 75 | 76 | stop_queries = torch.clamp(stop_queries, 0, query_limit) 77 | return self.x_final, stop_queries, dist, (dist <= self.epsilon) 78 | 79 | # check whether solution is found 80 | def search_succ(self, x, y, target, mask): 81 | self.queries[mask] += 1 82 | if target: 83 | return self.model.predict_label(x[mask]) == target[mask] 84 | else: 85 | return self.model.predict_label(x[mask]) != y[mask] 86 | 87 | # binary search for decision boundary along sgn direction 88 | def binary_search(self, x, y, target, sgn, valid_mask, tol=1e-3): 89 | sgn_norm = torch.norm(sgn.view(len(x), -1), 2, 1) 90 | sgn_unit = sgn / sgn_norm.view(len(x), 1, 1, 1) 91 | 92 | d_start = torch.zeros_like(y).float().cuda() 93 | d_end = self.d_t.clone() 94 | 95 | initial_succ_mask = self.search_succ(self.get_xadv(x, sgn_unit, self.d_t), y, target, valid_mask) 96 | to_search_ind = valid_mask.nonzero().flatten()[initial_succ_mask] 97 | d_end[to_search_ind] = torch.min(self.d_t, sgn_norm)[to_search_ind] 98 | 99 | while len(to_search_ind) > 0: 100 | d_mid = (d_start + d_end) / 2.0 101 | search_succ_mask = self.search_succ(self.get_xadv(x, sgn_unit, d_mid), y, target, to_search_ind) 102 | d_end[to_search_ind[search_succ_mask]] = d_mid[to_search_ind[search_succ_mask]] 103 | d_start[to_search_ind[~search_succ_mask]] = d_mid[to_search_ind[~search_succ_mask]] 104 | to_search_ind = to_search_ind[((d_end - d_start)[to_search_ind] > tol)] 105 | 106 | to_update_ind = (d_end < self.d_t).nonzero().flatten() 107 | if len(to_update_ind) > 0: 108 | self.d_t[to_update_ind] = d_end[to_update_ind] 109 | self.x_final[to_update_ind] = self.get_xadv(x, sgn_unit, d_end)[to_update_ind] 110 | self.sgn_t[to_update_ind] = sgn[to_update_ind] 111 | 112 | def __call__(self, data, label, target=None, query_limit=10000): 113 | return self.attack_hard_label(data, label, target=target, query_limit=query_limit) 114 | -------------------------------------------------------------------------------- /RayS_Single.py: -------------------------------------------------------------------------------- 1 | import time 2 | import numpy as np 3 | import torch 4 | import torch.nn.functional as F 5 | 6 | 7 | class RayS(object): 8 | def __init__(self, model, order=np.inf, epsilon=0.3, early_stopping=True): 9 | self.model = model 10 | self.order = order 11 | self.epsilon = epsilon 12 | self.sgn_t = None 13 | self.d_t = None 14 | self.x_final = None 15 | self.lin_search_rad = 10 16 | self.pre_set = {1, -1} 17 | self.early_stopping = early_stopping 18 | 19 | def get_xadv(self, x, v, d, lb=0., rb=1.): 20 | out = x + d * v 21 | return torch.clamp(out, lb, rb) 22 | 23 | def attack_hard_label(self, x, y, target=None, query_limit=10000, seed=None): 24 | """ Attack the original image and return adversarial example 25 | model: (pytorch model) 26 | (x, y): original image 27 | """ 28 | shape = list(x.shape) 29 | dim = np.prod(shape[1:]) 30 | if seed is not None: 31 | np.random.seed(seed) 32 | 33 | self.queries = 0 34 | self.d_t = np.inf 35 | self.sgn_t = torch.sign(torch.ones(shape)).cuda() 36 | self.x_final = self.get_xadv(x, self.sgn_t, self.d_t) 37 | dist = torch.tensor(np.inf) 38 | block_level = 0 39 | block_ind = 0 40 | 41 | for i in range(query_limit): 42 | 43 | block_num = 2 ** block_level 44 | block_size = int(np.ceil(dim / block_num)) 45 | start, end = block_ind * block_size, min(dim, (block_ind + 1) * block_size) 46 | 47 | attempt = self.sgn_t.clone().view(shape[0], dim) 48 | attempt[:, start:end] *= -1. 49 | attempt = attempt.view(shape) 50 | 51 | self.binary_search(x, y, target, attempt) 52 | 53 | block_ind += 1 54 | if block_ind == 2 ** block_level or end == dim: 55 | block_level += 1 56 | block_ind = 0 57 | 58 | dist = torch.norm(self.x_final - x, self.order) 59 | if self.early_stopping and (dist <= self.epsilon): 60 | break 61 | 62 | if self.queries >= query_limit: 63 | print('out of queries') 64 | break 65 | 66 | if i % 10 == 0: 67 | print("Iter %3d d_t %.8f dist %.8f queries %d" % (i + 1, self.d_t, dist, self.queries)) 68 | 69 | print("Iter %3d d_t %.6f dist %.6f queries %d" % (i + 1, self.d_t, dist, self.queries)) 70 | return self.x_final, self.queries, dist, (dist <= self.epsilon).float() 71 | 72 | def search_succ(self, x, y, target): 73 | self.queries += 1 74 | if target: 75 | return self.model.predict_label(x) == target 76 | else: 77 | return self.model.predict_label(x) != y 78 | 79 | def lin_search(self, x, y, target, sgn): 80 | d_end = np.inf 81 | for d in range(1, self.lin_search_rad + 1): 82 | if self.search_succ(self.get_xadv(x, sgn, d), y, target): 83 | d_end = d 84 | break 85 | return d_end 86 | 87 | def binary_search(self, x, y, target, sgn, tol=1e-3): 88 | sgn_unit = sgn / torch.norm(sgn) 89 | sgn_norm = torch.norm(sgn) 90 | 91 | d_start = 0 92 | if np.inf > self.d_t: # already have current result 93 | if not self.search_succ(self.get_xadv(x, sgn_unit, self.d_t), y, target): 94 | return False 95 | d_end = self.d_t 96 | else: # init run, try to find boundary distance 97 | d = self.lin_search(x, y, target, sgn) 98 | if d < np.inf: 99 | d_end = d * sgn_norm 100 | else: 101 | return False 102 | 103 | while (d_end - d_start) > tol: 104 | d_mid = (d_start + d_end) / 2.0 105 | if self.search_succ(self.get_xadv(x, sgn_unit, d_mid), y, target): 106 | d_end = d_mid 107 | else: 108 | d_start = d_mid 109 | if d_end < self.d_t: 110 | self.d_t = d_end 111 | self.x_final = self.get_xadv(x, sgn_unit, d_end) 112 | self.sgn_t = sgn 113 | return True 114 | else: 115 | return False 116 | 117 | def __call__(self, data, label, target=None, seed=None, query_limit=10000): 118 | return self.attack_hard_label(data, label, target=target, seed=seed, query_limit=query_limit) 119 | -------------------------------------------------------------------------------- /arch/cifar_model.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | 4 | class CIFAR10(nn.Module): 5 | def __init__(self): 6 | super(CIFAR10, self).__init__() 7 | self.features = self._make_layers() 8 | self.fc1 = nn.Linear(3200, 256) 9 | self.relu = nn.ReLU() 10 | self.fc2 = nn.Linear(256, 256) 11 | self.dropout = nn.Dropout(p=0.5) 12 | self.fc3 = nn.Linear(256, 10) 13 | 14 | def forward(self, x): 15 | out = self.features(x) 16 | out = out.view(out.size(0), -1) 17 | out = self.fc1(out) 18 | out = self.relu(out) 19 | out = self.dropout(out) 20 | out = self.fc2(out) 21 | out = self.relu(out) 22 | out = self.dropout(out) 23 | out = self.fc3(out) 24 | return out 25 | 26 | def _make_layers(self): 27 | layers = [] 28 | in_channels = 3 29 | layers += [nn.Conv2d(in_channels, 64, kernel_size=3), 30 | nn.BatchNorm2d(64), 31 | nn.ReLU()] 32 | layers += [nn.Conv2d(64, 64, kernel_size=3), 33 | nn.BatchNorm2d(64), 34 | nn.ReLU()] 35 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] 36 | layers += [nn.Conv2d(64, 128, kernel_size=3), 37 | nn.BatchNorm2d(128), 38 | nn.ReLU()] 39 | layers += [nn.Conv2d(128, 128, kernel_size=3), 40 | nn.BatchNorm2d(128), 41 | nn.ReLU()] 42 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] 43 | 44 | return nn.Sequential(*layers) 45 | -------------------------------------------------------------------------------- /arch/fs_utils.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | class Model_FS(nn.Module): 4 | def __init__(self, basic_net): 5 | super(Model_FS, self).__init__() 6 | self.basic_net = basic_net 7 | self.basic_net.eval() 8 | 9 | def forward(self, inputs): 10 | outputs, _ = self.basic_net(inputs) 11 | return outputs -------------------------------------------------------------------------------- /arch/madry_wrn.py: -------------------------------------------------------------------------------- 1 | # based on https://github.com/tensorflow/models/tree/master/resnet 2 | from __future__ import absolute_import 3 | from __future__ import division 4 | from __future__ import print_function 5 | 6 | import numpy as np 7 | import tensorflow as tf 8 | 9 | class Model(object): 10 | """ResNet model.""" 11 | 12 | def __init__(self, mode): 13 | """ResNet constructor. 14 | Args: 15 | mode: One of 'train' and 'eval'. 16 | """ 17 | self.mode = mode 18 | self._build_model() 19 | 20 | def add_internal_summaries(self): 21 | pass 22 | 23 | def _stride_arr(self, stride): 24 | """Map a stride scalar to the stride array for tf.nn.conv2d.""" 25 | return [1, stride, stride, 1] 26 | 27 | def _build_model(self): 28 | assert self.mode == 'train' or self.mode == 'eval' 29 | """Build the core model within the graph.""" 30 | with tf.variable_scope('input'): 31 | 32 | self.x_input = tf.placeholder( 33 | tf.float32, 34 | shape=[None, 32, 32, 3]) 35 | 36 | self.y_input = tf.placeholder(tf.int64, shape=None) 37 | 38 | 39 | input_standardized = tf.map_fn(lambda img: tf.image.per_image_standardization(img), 40 | self.x_input) 41 | x = self._conv('init_conv', input_standardized, 3, 3, 16, self._stride_arr(1)) 42 | 43 | 44 | 45 | strides = [1, 2, 2] 46 | activate_before_residual = [True, False, False] 47 | res_func = self._residual 48 | 49 | # wide residual network (https://arxiv.org/abs/1605.07146v1) 50 | # use filters = [16, 16, 32, 64] for a non-wide version 51 | filters = [16, 160, 320, 640] 52 | 53 | 54 | # Update hps.num_residual_units to 9 55 | 56 | with tf.variable_scope('unit_1_0'): 57 | x = res_func(x, filters[0], filters[1], self._stride_arr(strides[0]), 58 | activate_before_residual[0]) 59 | for i in range(1, 5): 60 | with tf.variable_scope('unit_1_%d' % i): 61 | x = res_func(x, filters[1], filters[1], self._stride_arr(1), False) 62 | 63 | with tf.variable_scope('unit_2_0'): 64 | x = res_func(x, filters[1], filters[2], self._stride_arr(strides[1]), 65 | activate_before_residual[1]) 66 | for i in range(1, 5): 67 | with tf.variable_scope('unit_2_%d' % i): 68 | x = res_func(x, filters[2], filters[2], self._stride_arr(1), False) 69 | 70 | with tf.variable_scope('unit_3_0'): 71 | x = res_func(x, filters[2], filters[3], self._stride_arr(strides[2]), 72 | activate_before_residual[2]) 73 | for i in range(1, 5): 74 | with tf.variable_scope('unit_3_%d' % i): 75 | x = res_func(x, filters[3], filters[3], self._stride_arr(1), False) 76 | 77 | with tf.variable_scope('unit_last'): 78 | x = self._batch_norm('final_bn', x) 79 | x = self._relu(x, 0.1) 80 | x = self._global_avg_pool(x) 81 | 82 | with tf.variable_scope('logit'): 83 | self.pre_softmax = self._fully_connected(x, 10) 84 | 85 | self.predictions = tf.argmax(self.pre_softmax, 1) 86 | self.correct_prediction = tf.equal(self.predictions, self.y_input) 87 | self.num_correct = tf.reduce_sum( 88 | tf.cast(self.correct_prediction, tf.int64)) 89 | self.accuracy = tf.reduce_mean( 90 | tf.cast(self.correct_prediction, tf.float32)) 91 | 92 | with tf.variable_scope('costs'): 93 | self.y_xent = tf.nn.sparse_softmax_cross_entropy_with_logits( 94 | logits=self.pre_softmax, labels=self.y_input) 95 | self.xent = tf.reduce_sum(self.y_xent, name='y_xent') 96 | self.mean_xent = tf.reduce_mean(self.y_xent) 97 | self.weight_decay_loss = self._decay() 98 | 99 | def _batch_norm(self, name, x): 100 | """Batch normalization.""" 101 | with tf.name_scope(name): 102 | return tf.contrib.layers.batch_norm( 103 | inputs=x, 104 | decay=.9, 105 | center=True, 106 | scale=True, 107 | activation_fn=None, 108 | updates_collections=None, 109 | is_training=(self.mode == 'train')) 110 | 111 | def _residual(self, x, in_filter, out_filter, stride, 112 | activate_before_residual=False): 113 | """Residual unit with 2 sub layers.""" 114 | if activate_before_residual: 115 | with tf.variable_scope('shared_activation'): 116 | x = self._batch_norm('init_bn', x) 117 | x = self._relu(x, 0.1) 118 | orig_x = x 119 | else: 120 | with tf.variable_scope('residual_only_activation'): 121 | orig_x = x 122 | x = self._batch_norm('init_bn', x) 123 | x = self._relu(x, 0.1) 124 | 125 | with tf.variable_scope('sub1'): 126 | x = self._conv('conv1', x, 3, in_filter, out_filter, stride) 127 | 128 | with tf.variable_scope('sub2'): 129 | x = self._batch_norm('bn2', x) 130 | x = self._relu(x, 0.1) 131 | x = self._conv('conv2', x, 3, out_filter, out_filter, [1, 1, 1, 1]) 132 | 133 | with tf.variable_scope('sub_add'): 134 | if in_filter != out_filter: 135 | orig_x = tf.nn.avg_pool(orig_x, stride, stride, 'VALID') 136 | orig_x = tf.pad( 137 | orig_x, [[0, 0], [0, 0], [0, 0], 138 | [(out_filter-in_filter)//2, (out_filter-in_filter)//2]]) 139 | x += orig_x 140 | 141 | tf.logging.debug('image after unit %s', x.get_shape()) 142 | return x 143 | 144 | def _decay(self): 145 | """L2 weight decay loss.""" 146 | costs = [] 147 | for var in tf.trainable_variables(): 148 | if var.op.name.find('DW') > 0: 149 | costs.append(tf.nn.l2_loss(var)) 150 | return tf.add_n(costs) 151 | 152 | def _conv(self, name, x, filter_size, in_filters, out_filters, strides): 153 | """Convolution.""" 154 | with tf.variable_scope(name): 155 | n = filter_size * filter_size * out_filters 156 | kernel = tf.get_variable( 157 | 'DW', [filter_size, filter_size, in_filters, out_filters], 158 | tf.float32, initializer=tf.random_normal_initializer( 159 | stddev=np.sqrt(2.0/n))) 160 | return tf.nn.conv2d(x, kernel, strides, padding='SAME') 161 | 162 | def _relu(self, x, leakiness=0.0): 163 | """Relu, with optional leaky support.""" 164 | return tf.where(tf.less(x, 0.0), leakiness * x, x, name='leaky_relu') 165 | 166 | def _fully_connected(self, x, out_dim): 167 | """FullyConnected layer for final output.""" 168 | num_non_batch_dimensions = len(x.shape) 169 | prod_non_batch_dimensions = 1 170 | for ii in range(num_non_batch_dimensions - 1): 171 | prod_non_batch_dimensions *= int(x.shape[ii + 1]) 172 | x = tf.reshape(x, [tf.shape(x)[0], -1]) 173 | w = tf.get_variable( 174 | 'DW', [prod_non_batch_dimensions, out_dim], 175 | initializer=tf.uniform_unit_scaling_initializer(factor=1.0)) 176 | b = tf.get_variable('biases', [out_dim], 177 | initializer=tf.constant_initializer()) 178 | return tf.nn.xw_plus_b(x, w, b) 179 | 180 | def _global_avg_pool(self, x): 181 | assert x.get_shape().ndims == 4 182 | return tf.reduce_mean(x, [1, 2]) 183 | -------------------------------------------------------------------------------- /arch/mnist_model.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | 4 | class MNIST(nn.Module): 5 | def __init__(self): 6 | super(MNIST, self).__init__() 7 | self.features = self._make_layers() 8 | self.fc1 = nn.Linear(1024, 200) 9 | self.relu = nn.ReLU() 10 | self.fc2 = nn.Linear(200, 200) 11 | self.dropout = nn.Dropout(p=0.5) 12 | self.fc3 = nn.Linear(200, 10) 13 | 14 | def forward(self, x): 15 | out = self.features(x) 16 | out = out.view(out.size(0), -1) 17 | out = self.fc1(out) 18 | out = self.relu(out) 19 | out = self.dropout(out) 20 | out = self.fc2(out) 21 | out = self.relu(out) 22 | out = self.dropout(out) 23 | out = self.fc3(out) 24 | return out 25 | 26 | def _make_layers(self): 27 | layers = [] 28 | in_channels = 1 29 | layers += [nn.Conv2d(in_channels, 32, kernel_size=3), 30 | nn.BatchNorm2d(32), 31 | nn.ReLU()] 32 | layers += [nn.Conv2d(32, 32, kernel_size=3), 33 | nn.BatchNorm2d(32), 34 | nn.ReLU()] 35 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] 36 | layers += [nn.Conv2d(32, 64, kernel_size=3), 37 | nn.BatchNorm2d(64), 38 | nn.ReLU()] 39 | layers += [nn.Conv2d(64, 64, kernel_size=3), 40 | nn.BatchNorm2d(64), 41 | nn.ReLU()] 42 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] 43 | 44 | return nn.Sequential(*layers) 45 | -------------------------------------------------------------------------------- /arch/preact_resnet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | 6 | class PreActBlock(nn.Module): 7 | '''Pre-activation version of the BasicBlock.''' 8 | expansion = 1 9 | 10 | def __init__(self, in_planes, planes, stride=1): 11 | super(PreActBlock, self).__init__() 12 | self.bn1 = nn.BatchNorm2d(in_planes) 13 | self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False) 14 | self.bn2 = nn.BatchNorm2d(planes) 15 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False) 16 | 17 | if stride != 1 or in_planes != self.expansion*planes: 18 | self.shortcut = nn.Sequential( 19 | nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False) 20 | ) 21 | 22 | def forward(self, x): 23 | out = F.relu(self.bn1(x)) 24 | shortcut = self.shortcut(x) if hasattr(self, 'shortcut') else x 25 | out = self.conv1(out) 26 | out = self.conv2(F.relu(self.bn2(out))) 27 | out += shortcut 28 | return out 29 | 30 | 31 | class PreActBottleneck(nn.Module): 32 | '''Pre-activation version of the original Bottleneck module.''' 33 | expansion = 4 34 | 35 | def __init__(self, in_planes, planes, stride=1): 36 | super(PreActBottleneck, self).__init__() 37 | self.bn1 = nn.BatchNorm2d(in_planes) 38 | self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False) 39 | self.bn2 = nn.BatchNorm2d(planes) 40 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False) 41 | self.bn3 = nn.BatchNorm2d(planes) 42 | self.conv3 = nn.Conv2d(planes, self.expansion*planes, kernel_size=1, bias=False) 43 | 44 | if stride != 1 or in_planes != self.expansion*planes: 45 | self.shortcut = nn.Sequential( 46 | nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False) 47 | ) 48 | 49 | def forward(self, x): 50 | out = F.relu(self.bn1(x)) 51 | shortcut = self.shortcut(out) if hasattr(self, 'shortcut') else x 52 | out = self.conv1(out) 53 | out = self.conv2(F.relu(self.bn2(out))) 54 | out = self.conv3(F.relu(self.bn3(out))) 55 | out += shortcut 56 | return out 57 | 58 | 59 | class PreActResNet(nn.Module): 60 | def __init__(self, block, num_blocks, num_classes=10): 61 | super(PreActResNet, self).__init__() 62 | self.in_planes = 64 63 | self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False) 64 | self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1) 65 | self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2) 66 | self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2) 67 | self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2) 68 | self.bn = nn.BatchNorm2d(512 * block.expansion) 69 | self.linear = nn.Linear(512 * block.expansion, num_classes) 70 | 71 | def _make_layer(self, block, planes, num_blocks, stride): 72 | strides = [stride] + [1]*(num_blocks-1) 73 | layers = [] 74 | for stride in strides: 75 | layers.append(block(self.in_planes, planes, stride)) 76 | self.in_planes = planes * block.expansion 77 | return nn.Sequential(*layers) 78 | 79 | def forward(self, x): 80 | out = self.conv1(x) 81 | out = self.layer1(out) 82 | out = self.layer2(out) 83 | out = self.layer3(out) 84 | out = self.layer4(out) 85 | out = F.relu(self.bn(out)) 86 | out = F.avg_pool2d(out, 4) 87 | out = out.view(out.size(0), -1) 88 | out = self.linear(out) 89 | return out 90 | 91 | 92 | def PreActResNet18(): 93 | return PreActResNet(PreActBlock, [2,2,2,2]) -------------------------------------------------------------------------------- /arch/wideresnet.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | 6 | 7 | class BasicBlock(nn.Module): 8 | def __init__(self, in_planes, out_planes, stride, dropRate=0.0): 9 | super(BasicBlock, self).__init__() 10 | self.bn1 = nn.BatchNorm2d(in_planes) 11 | self.relu1 = nn.ReLU(inplace=True) 12 | self.conv1 = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 13 | padding=1, bias=False) 14 | self.bn2 = nn.BatchNorm2d(out_planes) 15 | self.relu2 = nn.ReLU(inplace=True) 16 | self.conv2 = nn.Conv2d(out_planes, out_planes, kernel_size=3, stride=1, 17 | padding=1, bias=False) 18 | self.droprate = dropRate 19 | self.equalInOut = (in_planes == out_planes) 20 | self.convShortcut = (not self.equalInOut) and nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, 21 | padding=0, bias=False) or None 22 | 23 | def forward(self, x): 24 | if not self.equalInOut: 25 | x = self.relu1(self.bn1(x)) 26 | else: 27 | out = self.relu1(self.bn1(x)) 28 | out = self.relu2(self.bn2(self.conv1(out if self.equalInOut else x))) 29 | if self.droprate > 0: 30 | out = F.dropout(out, p=self.droprate, training=self.training) 31 | out = self.conv2(out) 32 | return torch.add(x if self.equalInOut else self.convShortcut(x), out) 33 | 34 | 35 | class NetworkBlock(nn.Module): 36 | def __init__(self, nb_layers, in_planes, out_planes, block, stride, dropRate=0.0): 37 | super(NetworkBlock, self).__init__() 38 | self.layer = self._make_layer(block, in_planes, out_planes, nb_layers, stride, dropRate) 39 | 40 | def _make_layer(self, block, in_planes, out_planes, nb_layers, stride, dropRate): 41 | layers = [] 42 | for i in range(int(nb_layers)): 43 | layers.append(block(i == 0 and in_planes or out_planes, out_planes, i == 0 and stride or 1, dropRate)) 44 | return nn.Sequential(*layers) 45 | 46 | def forward(self, x): 47 | return self.layer(x) 48 | 49 | 50 | class WideResNet(nn.Module): 51 | def __init__(self, depth=34, num_classes=10, widen_factor=10, dropRate=0.0): 52 | super(WideResNet, self).__init__() 53 | nChannels = [16, 16 * widen_factor, 32 * widen_factor, 64 * widen_factor] 54 | assert ((depth - 4) % 6 == 0) 55 | n = (depth - 4) / 6 56 | block = BasicBlock 57 | # 1st conv before any network block 58 | self.conv1 = nn.Conv2d(3, nChannels[0], kernel_size=3, stride=1, 59 | padding=1, bias=False) 60 | # 1st block 61 | self.block1 = NetworkBlock(n, nChannels[0], nChannels[1], block, 1, dropRate) 62 | # 1st sub-block 63 | self.sub_block1 = NetworkBlock(n, nChannels[0], nChannels[1], block, 1, dropRate) 64 | # 2nd block 65 | self.block2 = NetworkBlock(n, nChannels[1], nChannels[2], block, 2, dropRate) 66 | # 3rd block 67 | self.block3 = NetworkBlock(n, nChannels[2], nChannels[3], block, 2, dropRate) 68 | # global average pooling and classifier 69 | self.bn1 = nn.BatchNorm2d(nChannels[3]) 70 | self.relu = nn.ReLU(inplace=True) 71 | self.fc = nn.Linear(nChannels[3], num_classes) 72 | self.nChannels = nChannels[3] 73 | 74 | for m in self.modules(): 75 | if isinstance(m, nn.Conv2d): 76 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 77 | m.weight.data.normal_(0, math.sqrt(2. / n)) 78 | elif isinstance(m, nn.BatchNorm2d): 79 | m.weight.data.fill_(1) 80 | m.bias.data.zero_() 81 | elif isinstance(m, nn.Linear): 82 | m.bias.data.zero_() 83 | 84 | def forward(self, x): 85 | out = self.conv1(x) 86 | out = self.block1(out) 87 | out = self.block2(out) 88 | out = self.block3(out) 89 | out = self.relu(self.bn1(out)) 90 | out = F.avg_pool2d(out, 8) 91 | out = out.view(-1, self.nChannels) 92 | return self.fc(out) 93 | -------------------------------------------------------------------------------- /arch/wideresnet_compact.py: -------------------------------------------------------------------------------- 1 | 2 | import os 3 | import argparse 4 | import torch 5 | import torch.nn as nn 6 | import torchvision.datasets as datasets 7 | import torch.utils.data as data 8 | import torchvision.transforms as transforms 9 | import torch.nn.functional as F 10 | import math 11 | 12 | import sys 13 | 14 | class BasicBlock(nn.Module): 15 | def __init__(self, conv_layer, in_planes, out_planes, stride, dropRate=0.0): 16 | super(BasicBlock, self).__init__() 17 | self.bn1 = nn.BatchNorm2d(in_planes) 18 | self.relu1 = nn.ReLU(inplace=True) 19 | self.conv1 = conv_layer( 20 | in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False 21 | ) 22 | self.bn2 = nn.BatchNorm2d(out_planes) 23 | self.relu2 = nn.ReLU(inplace=True) 24 | self.conv2 = conv_layer( 25 | out_planes, out_planes, kernel_size=3, stride=1, padding=1, bias=False 26 | ) 27 | self.droprate = dropRate 28 | self.equalInOut = in_planes == out_planes 29 | self.convShortcut = ( 30 | (not self.equalInOut) 31 | and conv_layer( 32 | in_planes, 33 | out_planes, 34 | kernel_size=1, 35 | stride=stride, 36 | padding=0, 37 | bias=False, 38 | ) 39 | or None 40 | ) 41 | 42 | def forward(self, x): 43 | if not self.equalInOut: 44 | x = self.relu1(self.bn1(x)) 45 | else: 46 | out = self.relu1(self.bn1(x)) 47 | out = self.relu2(self.bn2(self.conv1(out if self.equalInOut else x))) 48 | if self.droprate > 0: 49 | out = F.dropout(out, p=self.droprate, training=self.training) 50 | out = self.conv2(out) 51 | return torch.add(x if self.equalInOut else self.convShortcut(x), out) 52 | 53 | 54 | class NetworkBlock(nn.Module): 55 | def __init__( 56 | self, nb_layers, in_planes, out_planes, block, conv_layer, stride, dropRate=0.0 57 | ): 58 | super(NetworkBlock, self).__init__() 59 | self.layer = self._make_layer( 60 | conv_layer, block, in_planes, out_planes, nb_layers, stride, dropRate 61 | ) 62 | 63 | def _make_layer( 64 | self, conv_layer, block, in_planes, out_planes, nb_layers, stride, dropRate 65 | ): 66 | layers = [] 67 | for i in range(int(nb_layers)): 68 | layers.append( 69 | block( 70 | conv_layer, 71 | i == 0 and in_planes or out_planes, 72 | out_planes, 73 | i == 0 and stride or 1, 74 | dropRate, 75 | ) 76 | ) 77 | return nn.Sequential(*layers) 78 | 79 | def forward(self, x): 80 | return self.layer(x) 81 | 82 | 83 | class WideResNet(nn.Module): 84 | def __init__( 85 | self, 86 | conv_layer, 87 | linear_layer, 88 | depth=34, 89 | num_classes=10, 90 | widen_factor=10, 91 | dropRate=0.0, 92 | ): 93 | super(WideResNet, self).__init__() 94 | nChannels = [16, 16 * widen_factor, 32 * widen_factor, 64 * widen_factor] 95 | assert (depth - 4) % 6 == 0 96 | n = (depth - 4) / 6 97 | block = BasicBlock 98 | # 1st conv before any network block 99 | self.conv1 = conv_layer( 100 | 3, nChannels[0], kernel_size=3, stride=1, padding=1, bias=False 101 | ) 102 | # 1st block 103 | self.block1 = NetworkBlock( 104 | n, nChannels[0], nChannels[1], block, conv_layer, 1, dropRate 105 | ) 106 | # 1st sub-block 107 | self.sub_block1 = NetworkBlock( 108 | n, nChannels[0], nChannels[1], block, conv_layer, 1, dropRate 109 | ) 110 | # 2nd block 111 | self.block2 = NetworkBlock( 112 | n, nChannels[1], nChannels[2], block, conv_layer, 2, dropRate 113 | ) 114 | # 3rd block 115 | self.block3 = NetworkBlock( 116 | n, nChannels[2], nChannels[3], block, conv_layer, 2, dropRate 117 | ) 118 | # global average pooling and classifier 119 | self.bn1 = nn.BatchNorm2d(nChannels[3]) 120 | self.relu = nn.ReLU(inplace=True) 121 | self.fc = linear_layer(nChannels[3], num_classes) 122 | self.nChannels = nChannels[3] 123 | 124 | for m in self.modules(): 125 | if isinstance(m, nn.Conv2d): 126 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 127 | m.weight.data.normal_(0, math.sqrt(2.0 / n)) 128 | elif isinstance(m, nn.BatchNorm2d): 129 | m.weight.data.fill_(1) 130 | m.bias.data.zero_() 131 | elif isinstance(m, linear_layer): 132 | m.bias.data.zero_() 133 | 134 | def forward(self, x): 135 | out = self.conv1(x) 136 | out = self.block1(out) 137 | out = self.block2(out) 138 | out = self.block3(out) 139 | out = self.relu(self.bn1(out)) 140 | out = F.avg_pool2d(out, 8) 141 | out = out.view(-1, self.nChannels) 142 | return self.fc(out) 143 | 144 | 145 | def wrn_28_10(): 146 | return WideResNet(nn.Conv2d, nn.Linear, depth=28, widen_factor=10) 147 | -------------------------------------------------------------------------------- /arch/wideresnet_fs.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | 6 | 7 | class BasicBlock(nn.Module): 8 | def __init__(self, in_planes, out_planes, stride, dropRate=0.0): 9 | super(BasicBlock, self).__init__() 10 | self.bn1 = nn.BatchNorm2d(in_planes) 11 | self.relu1 = nn.ReLU(inplace=True) 12 | self.conv1 = nn.Conv2d(in_planes, 13 | out_planes, 14 | kernel_size=3, 15 | stride=stride, 16 | padding=1, 17 | bias=False) 18 | self.bn2 = nn.BatchNorm2d(out_planes) 19 | self.relu2 = nn.ReLU(inplace=True) 20 | self.conv2 = nn.Conv2d(out_planes, 21 | out_planes, 22 | kernel_size=3, 23 | stride=1, 24 | padding=1, 25 | bias=False) 26 | self.droprate = dropRate 27 | self.equalInOut = (in_planes == out_planes) 28 | self.convShortcut = (not self.equalInOut) and nn.Conv2d( 29 | in_planes, 30 | out_planes, 31 | kernel_size=1, 32 | stride=stride, 33 | padding=0, 34 | bias=False) or None 35 | 36 | def forward(self, x): 37 | if not self.equalInOut: 38 | x = self.relu1(self.bn1(x)) 39 | else: 40 | out = self.relu1(self.bn1(x)) 41 | out = self.relu2(self.bn2(self.conv1(out if self.equalInOut else x))) 42 | if self.droprate > 0: 43 | out = F.dropout(out, p=self.droprate, training=self.training) 44 | out = self.conv2(out) 45 | return torch.add(x if self.equalInOut else self.convShortcut(x), out) 46 | 47 | 48 | class NetworkBlock(nn.Module): 49 | def __init__(self, 50 | nb_layers, 51 | in_planes, 52 | out_planes, 53 | block, 54 | stride, 55 | dropRate=0.0): 56 | super(NetworkBlock, self).__init__() 57 | self.layer = self._make_layer(block, in_planes, out_planes, nb_layers, 58 | stride, dropRate) 59 | 60 | def _make_layer(self, block, in_planes, out_planes, nb_layers, stride, 61 | dropRate): 62 | layers = [] 63 | for i in range(int(nb_layers)): 64 | layers.append( 65 | block(i == 0 and in_planes or out_planes, out_planes, 66 | i == 0 and stride or 1, dropRate)) 67 | return nn.Sequential(*layers) 68 | 69 | def forward(self, x): 70 | return self.layer(x) 71 | 72 | 73 | class WideResNet(nn.Module): 74 | def __init__(self, depth, num_classes, widen_factor=1, dropRate=0.0): 75 | super(WideResNet, self).__init__() 76 | nChannels = [ 77 | 16, 16 * widen_factor, 32 * widen_factor, 64 * widen_factor 78 | ] 79 | assert ((depth - 4) % 6 == 0) 80 | n = (depth - 4) / 6 81 | block = BasicBlock 82 | self.conv1 = nn.Conv2d(3, 83 | nChannels[0], 84 | kernel_size=3, 85 | stride=1, 86 | padding=1, 87 | bias=False) 88 | # 1st block 89 | self.block1 = NetworkBlock(n, nChannels[0], nChannels[1], block, 1, 90 | dropRate) 91 | # 2nd block 92 | self.block2 = NetworkBlock(n, nChannels[1], nChannels[2], block, 2, 93 | dropRate) 94 | # 3rd block 95 | self.block3 = NetworkBlock(n, nChannels[2], nChannels[3], block, 2, 96 | dropRate) 97 | # global average pooling and classifier 98 | self.bn1 = nn.BatchNorm2d(nChannels[3]) 99 | self.relu = nn.ReLU(inplace=True) 100 | self.fc = nn.Linear(nChannels[3], num_classes) 101 | self.nChannels = nChannels[3] 102 | 103 | for m in self.modules(): 104 | if isinstance(m, nn.Conv2d): 105 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 106 | m.weight.data.normal_(0, math.sqrt(2. / n)) 107 | elif isinstance(m, nn.BatchNorm2d): 108 | m.weight.data.fill_(1) 109 | m.bias.data.zero_() 110 | elif isinstance(m, nn.Linear): 111 | m.bias.data.zero_() 112 | 113 | def forward(self, x): 114 | out1 = self.conv1(x) 115 | out2 = self.block1(out1) 116 | out3 = self.block2(out2) 117 | out4 = self.block3(out3) 118 | out5 = self.relu(self.bn1(out4)) 119 | out6 = F.avg_pool2d(out5, 8) 120 | out = out6.view(-1, self.nChannels) 121 | 122 | return self.fc(out), self.fc(out) 123 | -------------------------------------------------------------------------------- /arch/wideresnet_he.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | class BasicBlock(nn.Module): 6 | def __init__(self, in_planes, out_planes, stride, dropRate=0.0): 7 | super(BasicBlock, self).__init__() 8 | self.bn1 = nn.BatchNorm2d(in_planes) 9 | self.relu1 = nn.ReLU(inplace=True) 10 | self.conv1 = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 11 | padding=1, bias=False) 12 | self.bn2 = nn.BatchNorm2d(out_planes) 13 | self.relu2 = nn.ReLU(inplace=True) 14 | self.conv2 = nn.Conv2d(out_planes, out_planes, kernel_size=3, stride=1, 15 | padding=1, bias=False) 16 | self.droprate = dropRate 17 | self.equalInOut = (in_planes == out_planes) 18 | self.convShortcut = (not self.equalInOut) and nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, 19 | padding=0, bias=False) or None 20 | 21 | def forward(self, x): 22 | if not self.equalInOut: 23 | x = self.relu1(self.bn1(x)) 24 | else: 25 | out = self.relu1(self.bn1(x)) 26 | out = self.relu2(self.bn2(self.conv1(out if self.equalInOut else x))) 27 | if self.droprate > 0: 28 | out = F.dropout(out, p=self.droprate, training=self.training) 29 | out = self.conv2(out) 30 | return torch.add(x if self.equalInOut else self.convShortcut(x), out) 31 | 32 | 33 | class NetworkBlock(nn.Module): 34 | def __init__(self, nb_layers, in_planes, out_planes, block, stride, dropRate=0.0): 35 | super(NetworkBlock, self).__init__() 36 | self.layer = self._make_layer(block, in_planes, out_planes, nb_layers, stride, dropRate) 37 | 38 | def _make_layer(self, block, in_planes, out_planes, nb_layers, stride, dropRate): 39 | layers = [] 40 | for i in range(int(nb_layers)): 41 | layers.append(block(i == 0 and in_planes or out_planes, out_planes, i == 0 and stride or 1, dropRate)) 42 | return nn.Sequential(*layers) 43 | 44 | def forward(self, x): 45 | return self.layer(x) 46 | 47 | 48 | class WideResNet(nn.Module): 49 | def __init__(self, depth=34, num_classes=10, widen_factor=10, dropRate=0.0, normalize = False): 50 | super(WideResNet, self).__init__() 51 | nChannels = [16, 16 * widen_factor, 32 * widen_factor, 64 * widen_factor] 52 | assert ((depth - 4) % 6 == 0) 53 | n = (depth - 4) / 6 54 | block = BasicBlock 55 | self.normalize = normalize 56 | # 1st conv before any network block 57 | self.conv1 = nn.Conv2d(3, nChannels[0], kernel_size=3, stride=1, 58 | padding=1, bias=False) 59 | # 1st block 60 | self.block1 = NetworkBlock(n, nChannels[0], nChannels[1], block, 1, dropRate) 61 | # 1st sub-block 62 | self.sub_block1 = NetworkBlock(n, nChannels[0], nChannels[1], block, 1, dropRate) 63 | # 2nd block 64 | self.block2 = NetworkBlock(n, nChannels[1], nChannels[2], block, 2, dropRate) 65 | # 3rd block 66 | self.block3 = NetworkBlock(n, nChannels[2], nChannels[3], block, 2, dropRate) 67 | # global average pooling and classifier 68 | self.bn1 = nn.BatchNorm2d(nChannels[3]) 69 | self.relu = nn.ReLU(inplace=True) 70 | if self.normalize: 71 | self.fc = nn.Linear(nChannels[3], num_classes, bias = False) 72 | else: 73 | self.fc = nn.Linear(nChannels[3], num_classes) 74 | self.nChannels = nChannels[3] 75 | 76 | for m in self.modules(): 77 | if isinstance(m, nn.Conv2d): 78 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 79 | m.weight.data.normal_(0, math.sqrt(2. / n)) 80 | elif isinstance(m, nn.BatchNorm2d): 81 | m.weight.data.fill_(1) 82 | m.bias.data.zero_() 83 | elif isinstance(m, nn.Linear) and not self.normalize: 84 | m.bias.data.zero_() 85 | 86 | def forward(self, x): 87 | out = self.conv1(x) 88 | out = self.block1(out) 89 | out = self.block2(out) 90 | out = self.block3(out) 91 | out = self.relu(self.bn1(out)) 92 | out = F.avg_pool2d(out, 8) 93 | out = out.view(-1, self.nChannels) 94 | if self.normalize: 95 | out = F.normalize(out, p=2, dim=1) 96 | for _, module in self.fc.named_modules(): 97 | if isinstance(module, nn.Linear): 98 | module.weight.data = F.normalize(module.weight, p=2, dim=1) 99 | return self.fc(out) -------------------------------------------------------------------------------- /arch/wideresnet_interp.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | 6 | 7 | class BasicBlock(nn.Module): 8 | def __init__(self, in_planes, out_planes, stride, dropRate=0.0): 9 | super(BasicBlock, self).__init__() 10 | self.bn1 = nn.BatchNorm2d(in_planes) 11 | self.relu1 = nn.ReLU(inplace=True) 12 | self.conv1 = nn.Conv2d(in_planes, 13 | out_planes, 14 | kernel_size=3, 15 | stride=stride, 16 | padding=1, 17 | bias=False) 18 | self.bn2 = nn.BatchNorm2d(out_planes) 19 | self.relu2 = nn.ReLU(inplace=True) 20 | self.conv2 = nn.Conv2d(out_planes, 21 | out_planes, 22 | kernel_size=3, 23 | stride=1, 24 | padding=1, 25 | bias=False) 26 | self.droprate = dropRate 27 | self.equalInOut = (in_planes == out_planes) 28 | self.convShortcut = (not self.equalInOut) and nn.Conv2d( 29 | in_planes, 30 | out_planes, 31 | kernel_size=1, 32 | stride=stride, 33 | padding=0, 34 | bias=False) or None 35 | 36 | def forward(self, x): 37 | if not self.equalInOut: 38 | x = self.relu1(self.bn1(x)) 39 | else: 40 | out = self.relu1(self.bn1(x)) 41 | out = self.relu2(self.bn2(self.conv1(out if self.equalInOut else x))) 42 | if self.droprate > 0: 43 | out = F.dropout(out, p=self.droprate, training=self.training) 44 | out = self.conv2(out) 45 | return torch.add(x if self.equalInOut else self.convShortcut(x), out) 46 | 47 | 48 | class NetworkBlock(nn.Module): 49 | def __init__(self, 50 | nb_layers, 51 | in_planes, 52 | out_planes, 53 | block, 54 | stride, 55 | dropRate=0.0): 56 | super(NetworkBlock, self).__init__() 57 | self.layer = self._make_layer(block, in_planes, out_planes, nb_layers, 58 | stride, dropRate) 59 | 60 | def _make_layer(self, block, in_planes, out_planes, nb_layers, stride, 61 | dropRate): 62 | layers = [] 63 | for i in range(int(nb_layers)): 64 | layers.append( 65 | block(i == 0 and in_planes or out_planes, out_planes, 66 | i == 0 and stride or 1, dropRate)) 67 | return nn.Sequential(*layers) 68 | 69 | def forward(self, x): 70 | return self.layer(x) 71 | 72 | 73 | class WideResNet(nn.Module): 74 | def __init__(self, depth, num_classes, widen_factor=1, dropRate=0.0): 75 | super(WideResNet, self).__init__() 76 | nChannels = [ 77 | 16, 16 * widen_factor, 32 * widen_factor, 64 * widen_factor 78 | ] 79 | assert ((depth - 4) % 6 == 0) 80 | n = (depth - 4) / 6 81 | block = BasicBlock 82 | self.conv1 = nn.Conv2d(3, 83 | nChannels[0], 84 | kernel_size=3, 85 | stride=1, 86 | padding=1, 87 | bias=False) 88 | # block 1 89 | self.block1 = NetworkBlock(n, nChannels[0], nChannels[1], block, 1, 90 | dropRate) 91 | # block 2 92 | self.block2 = NetworkBlock(n, nChannels[1], nChannels[2], block, 2, 93 | dropRate) 94 | # block 3 95 | self.block3 = NetworkBlock(n, nChannels[2], nChannels[3], block, 2, 96 | dropRate) 97 | 98 | self.bn1 = nn.BatchNorm2d(nChannels[3]) 99 | self.relu = nn.ReLU(inplace=True) 100 | self.fc = nn.Linear(nChannels[3], num_classes) 101 | self.nChannels = nChannels[3] 102 | 103 | for m in self.modules(): 104 | if isinstance(m, nn.Conv2d): 105 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 106 | m.weight.data.normal_(0, math.sqrt(2. / n)) 107 | elif isinstance(m, nn.BatchNorm2d): 108 | m.weight.data.fill_(1) 109 | m.bias.data.zero_() 110 | elif isinstance(m, nn.Linear): 111 | m.bias.data.zero_() 112 | 113 | def forward(self, x, mode='logits'): 114 | out = self.conv1(x) 115 | out = self.block1(out) 116 | out = self.block2(out) 117 | out = self.block3(out) 118 | out = self.relu(self.bn1(out)) 119 | out = F.avg_pool2d(out, 8) 120 | out = out.view(-1, self.nChannels) 121 | 122 | if mode.lower() == 'logits': 123 | return self.fc(out) 124 | elif mode.lower() == 'feature': 125 | return out.view(x.size(0), -1) 126 | else: 127 | raise Exception('unsupported mode is specified') -------------------------------------------------------------------------------- /arch/wideresnet_overfitting.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | 6 | 7 | class BasicBlock(nn.Module): 8 | def __init__(self, in_planes, out_planes, stride, dropRate=0.0): 9 | super(BasicBlock, self).__init__() 10 | self.bn1 = nn.BatchNorm2d(in_planes) 11 | self.relu1 = nn.ReLU(inplace=True) 12 | self.conv1 = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 13 | padding=1, bias=False) 14 | self.bn2 = nn.BatchNorm2d(out_planes) 15 | self.relu2 = nn.ReLU(inplace=True) 16 | self.conv2 = nn.Conv2d(out_planes, out_planes, kernel_size=3, stride=1, 17 | padding=1, bias=False) 18 | self.droprate = dropRate 19 | self.equalInOut = (in_planes == out_planes) 20 | self.convShortcut = (not self.equalInOut) and nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, 21 | padding=0, bias=False) or None 22 | def forward(self, x): 23 | if not self.equalInOut: 24 | x = self.relu1(self.bn1(x)) 25 | else: 26 | out = self.relu1(self.bn1(x)) 27 | out = self.relu2(self.bn2(self.conv1(out if self.equalInOut else x))) 28 | if self.droprate > 0: 29 | out = F.dropout(out, p=self.droprate, training=self.training) 30 | out = self.conv2(out) 31 | return torch.add(x if self.equalInOut else self.convShortcut(x), out) 32 | 33 | class NetworkBlock(nn.Module): 34 | def __init__(self, nb_layers, in_planes, out_planes, block, stride, dropRate=0.0): 35 | super(NetworkBlock, self).__init__() 36 | self.layer = self._make_layer(block, in_planes, out_planes, nb_layers, stride, dropRate) 37 | def _make_layer(self, block, in_planes, out_planes, nb_layers, stride, dropRate): 38 | layers = [] 39 | for i in range(int(nb_layers)): 40 | layers.append(block(i == 0 and in_planes or out_planes, out_planes, i == 0 and stride or 1, dropRate)) 41 | return nn.Sequential(*layers) 42 | def forward(self, x): 43 | return self.layer(x) 44 | 45 | class WideResNet(nn.Module): 46 | def __init__(self, depth, num_classes, widen_factor=1, dropRate=0.0): 47 | super(WideResNet, self).__init__() 48 | nChannels = [16, 16*widen_factor, 32*widen_factor, 64*widen_factor] 49 | assert((depth - 4) % 6 == 0) 50 | n = (depth - 4) / 6 51 | block = BasicBlock 52 | # 1st conv before any network block 53 | self.conv1 = nn.Conv2d(3, nChannels[0], kernel_size=3, stride=1, 54 | padding=1, bias=False) 55 | # 1st block 56 | self.block1 = NetworkBlock(n, nChannels[0], nChannels[1], block, 1, dropRate) 57 | # 2nd block 58 | self.block2 = NetworkBlock(n, nChannels[1], nChannels[2], block, 2, dropRate) 59 | # 3rd block 60 | self.block3 = NetworkBlock(n, nChannels[2], nChannels[3], block, 2, dropRate) 61 | # global average pooling and classifier 62 | self.bn1 = nn.BatchNorm2d(nChannels[3]) 63 | self.relu = nn.ReLU(inplace=True) 64 | self.fc = nn.Linear(nChannels[3], num_classes) 65 | self.nChannels = nChannels[3] 66 | 67 | for m in self.modules(): 68 | if isinstance(m, nn.Conv2d): 69 | nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 70 | elif isinstance(m, nn.BatchNorm2d): 71 | m.weight.data.fill_(1) 72 | m.bias.data.zero_() 73 | elif isinstance(m, nn.Linear): 74 | m.bias.data.zero_() 75 | def forward(self, x): 76 | out = self.conv1(x) 77 | out = self.block1(out) 78 | out = self.block2(out) 79 | out = self.block3(out) 80 | out = self.relu(self.bn1(out)) 81 | out = F.avg_pool2d(out, 8) 82 | out = out.view(-1, self.nChannels) 83 | return self.fc(out) -------------------------------------------------------------------------------- /arch/wideresnet_rst.py: -------------------------------------------------------------------------------- 1 | """Based on code from https://github.com/yaodongyu/TRADES""" 2 | 3 | import math 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | 8 | 9 | class BasicBlock(nn.Module): 10 | def __init__(self, in_planes, out_planes, stride, dropRate=0.0): 11 | super(BasicBlock, self).__init__() 12 | self.bn1 = nn.BatchNorm2d(in_planes) 13 | self.relu1 = nn.ReLU(inplace=True) 14 | self.conv1 = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 15 | padding=1, bias=False) 16 | self.bn2 = nn.BatchNorm2d(out_planes) 17 | self.relu2 = nn.ReLU(inplace=True) 18 | self.conv2 = nn.Conv2d(out_planes, out_planes, kernel_size=3, stride=1, 19 | padding=1, bias=False) 20 | self.droprate = dropRate 21 | self.equalInOut = (in_planes == out_planes) 22 | self.convShortcut = (not self.equalInOut) and nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, 23 | padding=0, bias=False) or None 24 | 25 | def forward(self, x): 26 | if not self.equalInOut: 27 | x = self.relu1(self.bn1(x)) 28 | else: 29 | out = self.relu1(self.bn1(x)) 30 | out = self.relu2(self.bn2(self.conv1(out if self.equalInOut else x))) 31 | if self.droprate > 0: 32 | out = F.dropout(out, p=self.droprate, training=self.training) 33 | out = self.conv2(out) 34 | return torch.add(x if self.equalInOut else self.convShortcut(x), out) 35 | 36 | 37 | class NetworkBlock(nn.Module): 38 | def __init__(self, nb_layers, in_planes, out_planes, block, stride, dropRate=0.0): 39 | super(NetworkBlock, self).__init__() 40 | self.layer = self._make_layer(block, in_planes, out_planes, nb_layers, stride, dropRate) 41 | 42 | def _make_layer(self, block, in_planes, out_planes, nb_layers, stride, dropRate): 43 | layers = [] 44 | for i in range(int(nb_layers)): 45 | layers.append(block(i == 0 and in_planes or out_planes, out_planes, i == 0 and stride or 1, dropRate)) 46 | return nn.Sequential(*layers) 47 | 48 | def forward(self, x): 49 | return self.layer(x) 50 | 51 | 52 | class WideResNet_RST(nn.Module): 53 | def __init__(self, depth=28, num_classes=10, widen_factor=10, dropRate=0.0): 54 | super(WideResNet_RST, self).__init__() 55 | nChannels = [16, 16 * widen_factor, 32 * widen_factor, 64 * widen_factor] 56 | assert ((depth - 4) % 6 == 0) 57 | n = (depth - 4) / 6 58 | block = BasicBlock 59 | # 1st conv before any network block 60 | self.conv1 = nn.Conv2d(3, nChannels[0], kernel_size=3, stride=1, 61 | padding=1, bias=False) 62 | # 1st block 63 | self.block1 = NetworkBlock(n, nChannels[0], nChannels[1], block, 1, dropRate) 64 | # 1st sub-block 65 | self.sub_block1 = NetworkBlock(n, nChannels[0], nChannels[1], block, 1, dropRate) 66 | # 2nd block 67 | self.block2 = NetworkBlock(n, nChannels[1], nChannels[2], block, 2, dropRate) 68 | # 3rd block 69 | self.block3 = NetworkBlock(n, nChannels[2], nChannels[3], block, 2, dropRate) 70 | # global average pooling and classifier 71 | self.bn1 = nn.BatchNorm2d(nChannels[3]) 72 | self.relu = nn.ReLU(inplace=True) 73 | self.fc = nn.Linear(nChannels[3], num_classes) 74 | self.nChannels = nChannels[3] 75 | 76 | for m in self.modules(): 77 | if isinstance(m, nn.Conv2d): 78 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 79 | m.weight.data.normal_(0, math.sqrt(2. / n)) 80 | elif isinstance(m, nn.BatchNorm2d): 81 | m.weight.data.fill_(1) 82 | m.bias.data.zero_() 83 | elif isinstance(m, nn.Linear): 84 | m.bias.data.zero_() 85 | 86 | def forward(self, x, return_prelogit=False): 87 | out = self.conv1(x) 88 | out = self.block1(out) 89 | out = self.block2(out) 90 | out = self.block3(out) 91 | out = self.relu(self.bn1(out)) 92 | out = F.avg_pool2d(out, 8) 93 | out = out.view(-1, self.nChannels) 94 | if return_prelogit: 95 | return self.fc(out), out 96 | else: 97 | return self.fc(out) -------------------------------------------------------------------------------- /attack_natural.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import json 3 | import numpy as np 4 | import torch 5 | import torchvision.models as models 6 | 7 | from dataset import load_mnist_test_data, load_cifar10_test_data, load_imagenet_test_data 8 | from general_torch_model import GeneralTorchModel 9 | 10 | from arch import mnist_model 11 | from arch import cifar_model 12 | 13 | from RayS_Single import RayS 14 | 15 | 16 | 17 | 18 | def main(): 19 | parser = argparse.ArgumentParser(description='Hard Label Attacks') 20 | parser.add_argument('--dataset', default='imagenet', type=str, 21 | help='Dataset') 22 | parser.add_argument('--targeted', default='0', type=str, 23 | help='targeted or untargeted') 24 | parser.add_argument('--norm', default='linf', type=str, 25 | help='Norm for attack, linf only') 26 | parser.add_argument('--num', default=10000, type=int, 27 | help='Number of samples to be attacked from test dataset.') 28 | parser.add_argument('--query', default=10000, type=int, 29 | help='Maximum queries for the attack') 30 | parser.add_argument('--batch', default=1, type=int, 31 | help='attack batch size.') 32 | parser.add_argument('--epsilon', default=0.05, type=float, 33 | help='attack strength') 34 | parser.add_argument('--early', default='1', type=str, 35 | help='early stopping (stop attack once the adversarial example is found)') 36 | args = parser.parse_args() 37 | 38 | targeted = True if args.targeted == '1' else False 39 | early_stopping = False if args.early == '0' else True 40 | order = 2 if args.norm == 'l2' else np.inf 41 | 42 | print(args) 43 | 44 | if args.dataset == 'mnist': 45 | model = mnist_model.MNIST().cuda() 46 | model = torch.nn.DataParallel(model, device_ids=[0]) 47 | model.load_state_dict(torch.load('model/mnist_gpu.pt')) 48 | test_loader = load_mnist_test_data(args.batch) 49 | torch_model = GeneralTorchModel(model, n_class=10, im_mean=None, im_std=None) 50 | elif args.dataset == 'cifar': 51 | model = cifar_model.CIFAR10().cuda() 52 | model = torch.nn.DataParallel(model, device_ids=[0]) 53 | model.load_state_dict(torch.load('model/cifar10_gpu.pt')) 54 | test_loader = load_cifar10_test_data(args.batch) 55 | torch_model = GeneralTorchModel(model, n_class=10, im_mean=None, im_std=None) 56 | elif args.dataset == 'resnet': 57 | model = models.__dict__["resnet50"](pretrained=True).cuda() 58 | model = torch.nn.DataParallel(model, device_ids=[0]) 59 | test_loader = load_imagenet_test_data(args.batch) 60 | torch_model = GeneralTorchModel(model, n_class=1000, im_mean=[0.485, 0.456, 0.406], 61 | im_std=[0.229, 0.224, 0.225]) 62 | elif args.dataset == 'inception': 63 | model = models.__dict__["inception_v3"](pretrained=True).cuda() 64 | model = torch.nn.DataParallel(model, device_ids=[0]) 65 | test_loader = load_imagenet_test_data(args.batch) 66 | torch_model = GeneralTorchModel(model, n_class=1000, im_mean=[0.485, 0.456, 0.406], 67 | im_std=[0.229, 0.224, 0.225]) 68 | else: 69 | print("Invalid dataset") 70 | exit(1) 71 | 72 | 73 | attack = RayS(torch_model, order=order, epsilon=args.epsilon, early_stopping=early_stopping) 74 | 75 | stop_dists = [] 76 | stop_queries = [] 77 | asr = [] 78 | np.random.seed(0) 79 | seeds = np.random.randint(10000, size=10000) 80 | count = 0 81 | for i, (xi, yi) in enumerate(test_loader): 82 | xi, yi = xi.cuda(), yi.cuda() 83 | 84 | if count == args.num: 85 | break 86 | 87 | if torch_model.predict_label(xi) != yi: 88 | continue 89 | 90 | np.random.seed(seeds[i]) 91 | 92 | target = np.random.randint(torch_model.n_class) * torch.ones(yi.shape, 93 | dtype=torch.long).cuda() if targeted else None 94 | while target and torch.sum(target == yi) > 0: 95 | print('re-generate target label') 96 | target = np.random.randint(torch_model.n_class) * torch.ones(len(xi), dtype=torch.long).cuda() 97 | 98 | adv, queries, dist, succ = attack(xi, yi, target=target, seed=seeds[i], 99 | query_limit=args.query) 100 | # print(queries, dist, succ) 101 | if succ: 102 | stop_queries.append(queries) 103 | if dist.item() < np.inf: 104 | stop_dists.append(dist.item()) 105 | elif early_stopping == False: 106 | if dist.item() < np.inf: 107 | stop_dists.append(dist.item()) 108 | 109 | asr.append(succ.item()) 110 | 111 | count += 1 112 | 113 | print("index: {:4d} avg dist: {:.4f} avg queries: {:.4f} asr: {:.4f} \n" 114 | .format(i, 115 | np.mean(np.array(stop_dists)), 116 | np.mean(np.array(stop_queries)), 117 | np.mean(np.array(asr)) 118 | )) 119 | 120 | 121 | name = args.dataset + '_' + args.alg + '_' + args.norm + '_query' + str(args.query) + '_eps' + str( 122 | args.epsilon) + '_early' + args.early 123 | summary_txt = 'distortion: ' + str(np.mean(np.array(stop_dists))) + ' queries: ' + str( 124 | np.mean(np.array(stop_queries))) + ' succ rate: ' + str(np.mean(np.array(asr))) 125 | with open(name + '_summary' + '.txt', 'w') as f: 126 | json.dump(summary_txt, f) 127 | 128 | 129 | if __name__ == "__main__": 130 | main() 131 | -------------------------------------------------------------------------------- /attack_robust.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import json 3 | import numpy as np 4 | import torch 5 | 6 | 7 | from dataset import load_mnist_test_data, load_cifar10_test_data, load_imagenet_test_data 8 | from general_torch_model import GeneralTorchModel 9 | from general_tf_model import GeneralTFModel 10 | 11 | from arch import fs_utils 12 | from arch import wideresnet 13 | from arch import wideresnet_fs 14 | from arch import wideresnet_interp 15 | from arch import wideresnet_he 16 | from arch import wideresnet_rst 17 | from arch import wideresnet_overfitting 18 | from arch import madry_wrn 19 | from arch import preact_resnet 20 | from arch import wideresnet_compact 21 | 22 | from RayS import RayS 23 | 24 | import torch.backends.cudnn as cudnn 25 | cudnn.benchmark = True 26 | 27 | np.random.seed(1234) 28 | 29 | 30 | def main(): 31 | parser = argparse.ArgumentParser(description='RayS Attacks') 32 | parser.add_argument('--dataset', default='rob_cifar_trades', type=str, 33 | help='robust model / dataset') 34 | parser.add_argument('--targeted', default='0', type=str, 35 | help='targeted or untargeted') 36 | parser.add_argument('--norm', default='linf', type=str, 37 | help='Norm for attack, linf only') 38 | parser.add_argument('--num', default=10000, type=int, 39 | help='Number of samples to be attacked from test dataset.') 40 | parser.add_argument('--query', default=10000, type=int, 41 | help='Maximum queries for the attack') 42 | parser.add_argument('--batch', default=10, type=int, 43 | help='attack batch size.') 44 | parser.add_argument('--epsilon', default=0.05, type=float, 45 | help='attack strength') 46 | args = parser.parse_args() 47 | 48 | targeted = True if args.targeted == '1' else False 49 | order = 2 if args.norm == 'l2' else np.inf 50 | 51 | print(args) 52 | summary_all = '' 53 | 54 | if args.dataset == 'rob_cifar_trades': 55 | model = wideresnet.WideResNet().cuda() 56 | model = torch.nn.DataParallel(model) 57 | model.module.load_state_dict(torch.load('model/rob_cifar_trades.pt')) 58 | test_loader = load_cifar10_test_data(args.batch) 59 | torch_model = GeneralTorchModel( 60 | model, n_class=10, im_mean=None, im_std=None) 61 | elif args.dataset == 'rob_cifar_adv': 62 | model = wideresnet.WideResNet().cuda() 63 | model = torch.nn.DataParallel(model) 64 | model.load_state_dict(torch.load('model/rob_cifar_madry.pt')) 65 | test_loader = load_cifar10_test_data(args.batch) 66 | torch_model = GeneralTorchModel( 67 | model, n_class=10, im_mean=None, im_std=None) 68 | elif args.dataset == 'rob_cifar_madry': 69 | import tensorflow as tf 70 | model = madry_wrn.Model(mode='eval') 71 | saver = tf.train.Saver() 72 | sess = tf.Session() 73 | saver.restore(sess, tf.train.latest_checkpoint('model/madry')) 74 | test_loader = load_cifar10_test_data(args.batch) 75 | torch_model = GeneralTFModel( 76 | model.pre_softmax, model.x_input, sess, n_class=10, im_mean=None, im_std=None) 77 | elif args.dataset == 'rob_cifar_interp': 78 | model = wideresnet_interp.WideResNet( 79 | depth=28, num_classes=10, widen_factor=10).cuda() 80 | model = torch.nn.DataParallel(model) 81 | checkpoint = torch.load('model/rob_cifar_interp') 82 | model.load_state_dict(checkpoint['net']) 83 | test_loader = load_cifar10_test_data(args.batch) 84 | torch_model = GeneralTorchModel(model, n_class=10, im_mean=[0.5, 0.5, 0.5], 85 | im_std=[0.5, 0.5, 0.5]) 86 | elif args.dataset == 'rob_cifar_fs': 87 | basic_net = wideresnet_fs.WideResNet( 88 | depth=28, num_classes=10, widen_factor=10).cuda() 89 | basic_net = basic_net.cuda() 90 | model = fs_utils.Model_FS(basic_net) 91 | model = torch.nn.DataParallel(model) 92 | checkpoint = torch.load('model/rob_cifar_fs') 93 | model.load_state_dict(checkpoint['net']) 94 | test_loader = load_cifar10_test_data(args.batch) 95 | torch_model = GeneralTorchModel(model, n_class=10, im_mean=[0.5, 0.5, 0.5], 96 | im_std=[0.5, 0.5, 0.5]) 97 | elif args.dataset == 'rob_cifar_sense': 98 | model = wideresnet.WideResNet().cuda() 99 | model = torch.nn.DataParallel(model) 100 | model.load_state_dict(torch.load( 101 | 'model/SENSE_checkpoint300.dict')['state_dict']) 102 | test_loader = load_cifar10_test_data(args.batch) 103 | torch_model = GeneralTorchModel( 104 | model, n_class=10, im_mean=None, im_std=None) 105 | elif args.dataset == 'rob_cifar_rst': 106 | model = wideresnet_rst.WideResNet_RST() 107 | model = torch.nn.DataParallel(model).cuda() 108 | model.load_state_dict(torch.load( 109 | 'model/rst_adv.pt.ckpt')['state_dict']) 110 | test_loader = load_cifar10_test_data(args.batch) 111 | torch_model = GeneralTorchModel( 112 | model, n_class=10, im_mean=None, im_std=None) 113 | elif args.dataset == 'rob_cifar_mart': 114 | model = wideresnet_rst.WideResNet_RST().cuda() 115 | model = torch.nn.DataParallel(model) 116 | model.load_state_dict(torch.load( 117 | 'model/mart_unlabel.pt')['state_dict']) 118 | test_loader = load_cifar10_test_data(args.batch) 119 | torch_model = GeneralTorchModel( 120 | model, n_class=10, im_mean=None, im_std=None) 121 | elif args.dataset == 'rob_cifar_uat': 122 | import tensorflow_hub as hub 123 | import tensorflow as tf 124 | UAT_HUB_URL = ('./model/uat_model') 125 | model = hub.Module(UAT_HUB_URL) 126 | my_input = tf.placeholder(tf.float32, shape=[None, 32, 32, 3]) 127 | my_logits = model(dict(x=my_input, decay_rate=0.1, prefix='default')) 128 | sess = tf.Session() 129 | sess.run([tf.global_variables_initializer(), tf.tables_initializer()]) 130 | test_loader = load_cifar10_test_data(args.batch) 131 | torch_model = GeneralTFModel( 132 | my_logits, my_input, sess, n_class=10, im_mean=[125.3/255, 123.0/255, 113.9/255], im_std=[63.0/255, 62.1/255, 66.7/255]) 133 | elif args.dataset == 'rob_cifar_overfitting': 134 | model = wideresnet_overfitting.WideResNet(depth=34, num_classes=10, widen_factor=20).cuda() 135 | model = torch.nn.DataParallel(model) 136 | model.load_state_dict(torch.load('model/rob_cifar_overfitting.pth')) 137 | test_loader = load_cifar10_test_data(args.batch) 138 | torch_model = GeneralTorchModel( 139 | model, n_class=10, im_mean=[0.4914, 0.4822, 0.4465], im_std=[0.2471, 0.2435, 0.2616]) 140 | elif args.dataset == 'rob_cifar_pretrain': 141 | model = wideresnet_overfitting.WideResNet(depth=28, num_classes=10, widen_factor=10).cuda() 142 | model = torch.nn.DataParallel(model) 143 | model.load_state_dict(torch.load('model/rob_cifar_pretrain.pt')) 144 | test_loader = load_cifar10_test_data(args.batch) 145 | torch_model = GeneralTorchModel( 146 | model, n_class=10, im_mean=[0.5, 0.5, 0.5], im_std=[0.5, 0.5, 0.5]) 147 | elif args.dataset == 'rob_cifar_fast': 148 | model = preact_resnet.PreActResNet18().cuda() 149 | model.load_state_dict(torch.load('model/rob_cifar_fast_epoch30.pth')) 150 | test_loader = load_cifar10_test_data(args.batch) 151 | torch_model = GeneralTorchModel( 152 | model, n_class=10, im_mean=[0.4914, 0.4822, 0.4465], im_std=[0.2471, 0.2435, 0.2616]) 153 | elif args.dataset == 'rob_cifar_compact': 154 | model = torch.nn.DataParallel(wideresnet_compact.wrn_28_10()) 155 | ckpt = torch.load('model/rob_cifar_compact.pth.tar', map_location="cpu")["state_dict"] 156 | model.load_state_dict(ckpt) 157 | model.cuda() 158 | test_loader = load_cifar10_test_data(args.batch) 159 | torch_model = GeneralTorchModel( 160 | model, n_class=10, im_mean=None, im_std=None) 161 | if args.dataset == 'rob_cifar_mma': 162 | from advertorch_examples.models import get_cifar10_wrn28_widen_factor 163 | model = get_cifar10_wrn28_widen_factor(4).cuda() 164 | model = torch.nn.DataParallel(model) 165 | model.module.load_state_dict(torch.load('model/rob_cifar_mma.pt')['model']) 166 | test_loader = load_cifar10_test_data(args.batch) 167 | torch_model = GeneralTorchModel( 168 | model, n_class=10, im_mean=None, im_std=None) 169 | if args.dataset == 'rob_cifar_he': 170 | model = wideresnet_he.WideResNet(normalize = True).cuda() 171 | model = torch.nn.DataParallel(model) 172 | model.module.load_state_dict(torch.load('model/rob_cifar_pgdHE.pt')) 173 | test_loader = load_cifar10_test_data(args.batch) 174 | torch_model = GeneralTorchModel( 175 | model, n_class=10, im_mean=None, im_std=None) 176 | else: 177 | print("Invalid dataset") 178 | exit(1) 179 | 180 | attack = RayS(torch_model, epsilon=args.epsilon, order=order) 181 | 182 | adbd = [] 183 | queries = [] 184 | succ = [] 185 | 186 | count = 0 187 | for i, (data, label) in enumerate(test_loader): 188 | data, label = data.cuda(), label.cuda() 189 | 190 | if count >= args.num: 191 | break 192 | 193 | if targeted: 194 | target = np.random.randint(torch_model.n_class) * torch.ones( 195 | label.shape, dtype=torch.long).cuda() if targeted else None 196 | while target and torch.sum(target == label) > 0: 197 | print('re-generate target label') 198 | target = np.random.randint( 199 | torch_model.n_class) * torch.ones(len(data), dtype=torch.long).cuda() 200 | else: 201 | target = None 202 | 203 | _, queries_b, adbd_b, succ_b = attack( 204 | data, label, target=target, query_limit=args.query) 205 | 206 | queries.append(queries_b) 207 | adbd.append(adbd_b) 208 | succ.append(succ_b) 209 | 210 | count += data.shape[0] 211 | 212 | summary_batch = "Batch: {:4d} Avg Queries (when found adversarial examples): {:.4f} ADBD: {:.4f} Robust Acc: {:.4f}\n" \ 213 | .format( 214 | i + 1, 215 | torch.stack(queries).flatten().float().mean(), 216 | torch.stack(adbd).flatten().mean(), 217 | 1 - torch.stack(succ).flatten().float().mean() 218 | ) 219 | print(summary_batch) 220 | summary_all += summary_batch 221 | 222 | name = args.dataset + '_query_' + str(args.query) + '_batch' 223 | with open(name + '_summary' + '.txt', 'w') as fileopen: 224 | json.dump(summary_all, fileopen) 225 | 226 | 227 | if __name__ == "__main__": 228 | main() 229 | -------------------------------------------------------------------------------- /dataset.py: -------------------------------------------------------------------------------- 1 | import random 2 | import numpy as np 3 | import torch 4 | import torchvision.datasets as dsets 5 | import torchvision.transforms as transforms 6 | 7 | 8 | def load_mnist_test_data(test_batch_size=1): 9 | """ Load MNIST data from torchvision.datasets 10 | input: None 11 | output: minibatches of train and test sets 12 | """ 13 | # MNIST Dataset 14 | test_dataset = dsets.MNIST(root='./data/mnist', train=False, transform=transforms.ToTensor()) 15 | test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=test_batch_size, shuffle=False) 16 | 17 | return test_loader 18 | 19 | 20 | def load_cifar10_test_data(test_batch_size=1): 21 | # CIFAR10 Dataset 22 | test_dataset = dsets.CIFAR10('./data/cifar10-py', download=True, train=False, transform=transforms.ToTensor()) 23 | test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=test_batch_size, shuffle=False) 24 | 25 | return test_loader 26 | 27 | 28 | def load_imagenet_test_data(test_batch_size=1, folder='../val/'): 29 | val_dataset = dsets.ImageFolder( 30 | folder, 31 | transforms.Compose([ 32 | transforms.Resize(256), 33 | transforms.CenterCrop(224), 34 | transforms.ToTensor(), 35 | ])) 36 | 37 | rand_seed = 42 38 | 39 | torch.manual_seed(rand_seed) 40 | torch.cuda.manual_seed(rand_seed) 41 | np.random.seed(rand_seed) 42 | random.seed(rand_seed) 43 | torch.backends.cudnn.deterministic = True 44 | val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=test_batch_size, shuffle=True) 45 | 46 | return val_loader 47 | -------------------------------------------------------------------------------- /general_tf_model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import torch.nn as nn 4 | 5 | import tensorflow as tf 6 | # import tensorflow.compat.v1 as tf 7 | import tensorflow_hub as hub 8 | 9 | class GeneralTFModel(nn.Module): 10 | def __init__(self, model_logits, x_input, sess, n_class=10, im_mean=None, im_std=None): 11 | super(GeneralTFModel, self).__init__() 12 | self.model_logits = model_logits 13 | self.x_input = x_input 14 | self.sess = sess 15 | self.num_queries = 0 16 | self.im_mean = im_mean 17 | self.im_std = im_std 18 | self.n_class = n_class 19 | 20 | def forward(self, image): 21 | if len(image.size()) != 4: 22 | image = image.unsqueeze(0) 23 | image_tf = np.moveaxis(image.cpu().numpy(), 1, 3) 24 | logits = self.sess.run(self.model_logits, {self.x_input: image_tf}) 25 | return torch.from_numpy(logits).cuda() 26 | 27 | def preprocess(self, image): 28 | if isinstance(image, np.ndarray): 29 | processed = torch.from_numpy(image).type(torch.FloatTensor) 30 | else: 31 | processed = image 32 | 33 | if self.im_mean is not None and self.im_std is not None: 34 | im_mean = torch.tensor(self.im_mean).cuda().view(1, processed.shape[1], 1, 1).repeat( 35 | processed.shape[0], 1, 1, 1) 36 | im_std = torch.tensor(self.im_std).cuda().view(1, processed.shape[1], 1, 1).repeat( 37 | processed.shape[0], 1, 1, 1) 38 | processed = (processed - im_mean) / im_std 39 | return processed 40 | 41 | def predict_prob(self, image): 42 | if len(image.size()) != 4: 43 | image = image.unsqueeze(0) 44 | image = self.preprocess(image) 45 | self.num_queries += image.size(0) 46 | 47 | image_tf = np.moveaxis(image.cpu().numpy(), 1, 3) 48 | logits = self.sess.run(self.model_logits, {self.x_input: image_tf}) 49 | 50 | return torch.from_numpy(logits).cuda() 51 | 52 | def predict_label(self, image): 53 | logits = self.predict_prob(image) 54 | _, predict = torch.max(logits, 1) 55 | return predict 56 | -------------------------------------------------------------------------------- /general_torch_model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import torch.nn as nn 4 | 5 | 6 | class GeneralTorchModel(nn.Module): 7 | def __init__(self, model, n_class=10, im_mean=None, im_std=None): 8 | super(GeneralTorchModel, self).__init__() 9 | self.model = model 10 | self.model.eval() 11 | self.num_queries = 0 12 | self.im_mean = im_mean 13 | self.im_std = im_std 14 | self.n_class = n_class 15 | 16 | def forward(self, image): 17 | if len(image.size()) != 4: 18 | image = image.unsqueeze(0) 19 | image = self.preprocess(image) 20 | logits = self.model(image) 21 | return logits 22 | 23 | def preprocess(self, image): 24 | if isinstance(image, np.ndarray): 25 | processed = torch.from_numpy(image).type(torch.FloatTensor) 26 | else: 27 | processed = image 28 | 29 | if self.im_mean is not None and self.im_std is not None: 30 | im_mean = torch.tensor(self.im_mean).cuda().view(1, processed.shape[1], 1, 1).repeat( 31 | processed.shape[0], 1, 1, 1) 32 | im_std = torch.tensor(self.im_std).cuda().view(1, processed.shape[1], 1, 1).repeat( 33 | processed.shape[0], 1, 1, 1) 34 | processed = (processed - im_mean) / im_std 35 | return processed 36 | 37 | def predict_prob(self, image): 38 | with torch.no_grad(): 39 | if len(image.size()) != 4: 40 | image = image.unsqueeze(0) 41 | image = self.preprocess(image) 42 | logits = self.model(image) 43 | self.num_queries += image.size(0) 44 | return logits 45 | 46 | def predict_label(self, image): 47 | logits = self.predict_prob(image) 48 | _, predict = torch.max(logits, 1) 49 | return predict -------------------------------------------------------------------------------- /images/adbd_leaderboard.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/uclaml/RayS/29bc17493cde3ade6ca0027ae2c318785eceb5dd/images/adbd_leaderboard.png -------------------------------------------------------------------------------- /model/cifar10_gpu.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/uclaml/RayS/29bc17493cde3ade6ca0027ae2c318785eceb5dd/model/cifar10_gpu.pt -------------------------------------------------------------------------------- /model/mnist_gpu.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/uclaml/RayS/29bc17493cde3ade6ca0027ae2c318785eceb5dd/model/mnist_gpu.pt -------------------------------------------------------------------------------- /pgbar.py: -------------------------------------------------------------------------------- 1 | import os, torch 2 | import sys 3 | import time 4 | 5 | _, term_width = os.popen('stty size', 'r').read().split() 6 | term_width = int(term_width) 7 | TOTAL_BAR_LENGTH = 65. 8 | last_time = time.time() 9 | begin_time = last_time 10 | 11 | 12 | def progress_bar(current, total, msg=None): 13 | global last_time, begin_time 14 | if current == 0: 15 | begin_time = time.time() # Reset for new bar. 16 | 17 | cur_len = int(TOTAL_BAR_LENGTH * current / total) 18 | rest_len = int(TOTAL_BAR_LENGTH - cur_len) - 1 19 | 20 | sys.stdout.write(' [') 21 | for i in range(cur_len): 22 | sys.stdout.write('=') 23 | sys.stdout.write('>') 24 | for i in range(rest_len): 25 | sys.stdout.write('.') 26 | sys.stdout.write(']') 27 | 28 | cur_time = time.time() 29 | step_time = cur_time - last_time 30 | last_time = cur_time 31 | tot_time = cur_time - begin_time 32 | 33 | L = [] 34 | L.append(' Step: %s' % format_time(step_time)) 35 | L.append(' | Tot: %s' % format_time(tot_time)) 36 | if msg: 37 | L.append(' | ' + msg) 38 | 39 | msg = ''.join(L) 40 | sys.stdout.write(msg) 41 | for i in range(term_width - int(TOTAL_BAR_LENGTH) - len(msg) - 3): 42 | sys.stdout.write(' ') 43 | 44 | # Go back to the center of the bar. 45 | for i in range(term_width - int(TOTAL_BAR_LENGTH / 2) + 2): 46 | sys.stdout.write('\b') 47 | sys.stdout.write(' %d/%d ' % (current + 1, total)) 48 | 49 | if current < total - 1: 50 | sys.stdout.write('\r') 51 | else: 52 | sys.stdout.write('\n') 53 | sys.stdout.flush() 54 | 55 | 56 | def format_time(seconds): 57 | days = int(seconds / 3600 / 24) 58 | seconds = seconds - days * 3600 * 24 59 | hours = int(seconds / 3600) 60 | seconds = seconds - hours * 3600 61 | minutes = int(seconds / 60) 62 | seconds = seconds - minutes * 60 63 | secondsf = int(seconds) 64 | seconds = seconds - secondsf 65 | millis = int(seconds * 1000) 66 | 67 | f = '' 68 | i = 1 69 | if days > 0: 70 | f += str(days) + 'D' 71 | i += 1 72 | if hours > 0 and i <= 2: 73 | f += str(hours) + 'h' 74 | i += 1 75 | if minutes > 0 and i <= 2: 76 | f += str(minutes) + 'm' 77 | i += 1 78 | if secondsf > 0 and i <= 2: 79 | f += str(secondsf) + 's' 80 | i += 1 81 | if millis > 0 and i <= 2: 82 | f += str(millis) + 'ms' 83 | i += 1 84 | if f == '': 85 | f = '0ms' 86 | return f 87 | --------------------------------------------------------------------------------