def clone_g0_inputs_on_ngpus(self, inputs, outputs, g0_inputs): """ Clone variables unused by the attack on all GPUs. Specifically, the ground-truth label, y, has to be preserved until the training step. :param inputs: A list of dictionaries as the inputs to each step. :param outputs: A list of dictionaries as the outputs of each step. :param g0_inputs: Initial variables to be cloned. :return: Updated inputs and outputs. """ assert len(inputs) == len(outputs), ( 'Inputs and outputs should have the same number of elements.') inputs[0].update(g0_inputs) outputs[0].update(g0_inputs) # Copy g0_inputs forward for i in range(1, len(inputs)): # Create the graph for i'th step of attack device_name = inputs[i]['x'].device with tf.device(device_name): with tf.variable_scope('step%d' % i): for k, v in g0_inputs.iteritems(): if k not in inputs[i]: v_copy = clone_variable(k, v) inputs[i][k] = v_copy outputs[i][k] = v_copy return inputs, outputs
def _create_train_graph(self): super(TrainerMultiGPU, self)._create_train_graph() assert '_multigpu' in self.hparams.attack_type_train hparams = self.hparams model = self.model sess = self.sess # Create trainable variables on last gpu. # Variables are set to trainable or non-trainable first time they are # created. This caused a bug when the last gpu is used both for attack # generation and training. With this bug the result of naive training # was affected by the length of the unused adversarial generation # graph. device_name = '/gpu:%d' % (hparams.ngpu-1) model.set_device(device_name) with tf.device(device_name): x = clone_variable('x', self.g0_inputs['x']) model.set_training(training=True) preds = model.get_probs(x) # Generates steps on gpus model.set_training(training=False) logging.info("Initializing train attack %s" % hparams.attack_type_train) inputs, outputs = create_adv_by_name( model, self.g0_inputs['x'], hparams.attack_type_train, sess, y=self.g0_inputs['y'], nb_iter=hparams.attack_nb_iter_train, dataset=hparams.dataset, ngpu=hparams.ngpu) inputs, outputs = self.clone_g0_inputs_on_ngpus( inputs, outputs, self.g0_inputs) # Train step on last gpu device_name = '/gpu:%d' % (hparams.ngpu-1) model.set_device(device_name) with tf.device(device_name): with tf.variable_scope('last'): inputs += [OrderedDict()] for k, v in outputs[-1].iteritems(): v_copy = clone_variable(k, v) inputs[-1][k] = v_copy x = inputs[-1]['x'] adv_x = inputs[-1]['adv_x'] y = inputs[-1]['y'] if not hparams.adv_train: model.set_training(training=True) preds = model.get_probs(x) preds_adv = None elif not hparams.only_adv_train: model.set_training(training=True) preds = model.get_probs(x) model.set_training(training=True) preds_adv = model.get_probs(adv_x) else: preds = None model.set_training(training=True) preds_adv = model.get_probs(adv_x) train_fetches = self._build_train_op(preds, y, preds_adv) outputs += [{'fetches': train_fetches}] # Create the sync operation device_name = '/gpu:%d' % (hparams.ngpu-1) model.set_device(device_name) with tf.device(device_name): sync_ops = model.create_sync_ops(host_device=device_name) self.inputs = inputs self.outputs = outputs self.sync_ops = sync_ops
def attack(self, x, y_p, **kwargs): """ This method creates a symoblic graph of the MadryEtAl attack on multiple GPUs. The graph is created on the first n GPUs. Stop gradient is needed to get the speed-up. This prevents us from being able to back-prop through the attack. :param x: A tensor with the input image. :param y_p: Ground truth label or predicted label. :return: Two lists containing the input and output tensors of each GPU. """ inputs = [] outputs = [] # Create the initial random perturbation device_name = '/gpu:0' self.model.set_device(device_name) with tf.device(device_name): with tf.variable_scope('init_rand'): if self.rand_init: eta = tf.random_uniform(tf.shape(x), -self.eps, self.eps) eta = clip_eta(eta, self.ord, self.eps) eta = tf.stop_gradient(eta) else: eta = tf.zeros_like(x) # TODO: Break the graph only nGPU times instead of nb_iter times. # The current implementation by the time an adversarial example is # used for training, the weights of the model have changed nb_iter # times. This can cause slower convergence compared to the single GPU # adversarial training. for i in range(self.nb_iter): # Create the graph for i'th step of attack inputs += [OrderedDict()] outputs += [OrderedDict()] device_name = x.device self.model.set_device(device_name) with tf.device(device_name): with tf.variable_scope('step%d' % i): if i > 0: # Clone the variables to separate the graph of 2 GPUs x = clone_variable('x', x) y_p = clone_variable('y_p', y_p) eta = clone_variable('eta', eta) inputs[i]['x'] = x inputs[i]['y_p'] = y_p outputs[i]['x'] = x outputs[i]['y_p'] = y_p inputs[i]['eta'] = eta eta = self.attack_single_step(x, eta, y_p) if i < self.nb_iter-1: outputs[i]['eta'] = eta else: # adv_x, not eta is the output of the last step adv_x = x + eta if (self.clip_min is not None and self.clip_max is not None): adv_x = tf.clip_by_value(adv_x, self.clip_min, self.clip_max) adv_x = tf.stop_gradient(adv_x, name='adv_x') outputs[i]['adv_x'] = adv_x return inputs, outputs