//  [python]
hps.n_train, hps.n_test, hps.n_classes = dataset.n_train, dataset.n_test, dataset.n_classes
hps.height, hps.width, hps.n_colors = dataset.height, dataset.width, dataset.n_colors
channel_mean, channel_std = np.array([[[[0.485, 0.456, 0.406]]]]), np.array([[[[0.229, 0.224, 0.225]]]])

# Define the computational graph
graph = tf.Graph()
with graph.as_default(), tf.device('/gpu:0'):
    #tf.set_random_seed(1) #does not ensure reproducibility
    tf_n_updates = tf.Variable(0, trainable=False)

    flag_train = tf.placeholder(tf.bool, name='is_training')
    at_flag_tf = tf.placeholder(tf.bool, name='at_flag')
    rub_flag_tf = tf.placeholder(tf.bool, name='rub_flag')
    adv_flag_tf = tf.placeholder(tf.bool, name='adv_flag')
    max_conf_flag_tf = tf.placeholder(tf.bool, name='max_conf_flag')
    at_frac_tf = tf.placeholder(tf.float32, name='at_frac')
    pgd_niter_tf = tf.placeholder_with_default(0, shape=())  # if just placeholder is used - weird bug in tf.while
    pgd_stepsize_tf = tf.cast(hps.pgd_step, tf.float32)
    x_in = tf.placeholder(tf.float32, [None, hps.height, hps.width, hps.n_colors])
    y_in = tf.placeholder(tf.int32, [None, ])
    hps.n_ex = tf.shape(x_in)[0]

    tower_grads = []
    imgs_per_gpu = hps.batch_size // len(hps.gpus)
    losses, regs, loss_upds, reg_upds = [], [], [], []
    acc_rates, avg_confs, acc_rate_upds, avg_conf_upds = [], [], [], []
    loss_clean_metric, loss_clean_metric_upd, loss_out_metric, loss_out_metric_upd, loss_total_metric, loss_total_metric_upd, reg_metric, reg_metric_upd = [], [], [], [], [], [], [], []
    losses_clean_metric, losses_clean_metric_upds, losses_out_metric, losses_out_metric_upds, losses_total_metric, losses_total_metric_upds, regs_metric, regs_metric_upds = [], [], [], [], [], [], [], []    
    logits_list, probs_list, labels_list = [], [], []
    with tf.variable_scope(tf.get_variable_scope()):
        for i in range(len(hps.gpus)):
            with tf.device('/gpu:%d' % i), tf.name_scope('tower_%d' % i) as scope:
                id_from, id_to = i * imgs_per_gpu, i * imgs_per_gpu + imgs_per_gpu
                x, y = x_in[id_from:id_to], y_in[id_from:id_to]
                y = tf.one_hot(y, hps.n_classes)  # needed for rubbish adversarial training
                print('Batch on gpu {}: from {} to {}, {}'.format(i, id_from, id_to, x))

                model = models.models_dict[hps.model](hps)
                n_clean = tf.cast(batch_size_per_gpu, tf.int32)
                n_adv = tf.cast(at_frac_tf * batch_size_per_gpu, tf.int32)
                x_adv, y_adv = utils_tf.gen_adv(model, tf.concat([x, x, x], axis=0)[:n_adv], tf.concat([y, y, y], axis=0)[:n_adv], hps.p, hps.pgd_eps, pgd_niter_tf,
                                                pgd_stepsize_tf, rub_flag_tf, adv_flag_tf, hps.frac_perm,
                                                hps.lowpass, max_conf_flag_tf)
                # During training or if x_adv is empty: use clean and adversarial/rubbish examples.
                # During evaluation: use adversarial/rubbish examples only (which may still be clean iff
                # rub_flag_tf: False and adv_flag_tf: False).
                x_c, y_c = tf.cond(tf.logical_or(flag_train, tf.equal(tf.shape(x_adv)[0], 0)),
                               lambda: (tf.concat([x_adv, x], axis=0), tf.concat([y_adv, y], axis=0)),
                               lambda: (x_adv, y_adv))

                # Reuse the model declared for adv. examples here, and also for the next tower.
                tf.get_variable_scope().reuse_variables()

                logits_c = model.get_logits(x_c, flag_train=flag_train)
                probs_c = tf.nn.softmax(logits_c)
                logits_adv =  logits_c[:n_adv]
                logits = logits_c[n_adv:]
                maxclass = tf.argmax(logits_adv, axis=-1)
                loss_adv = tf.cond(max_conf_flag_tf,
                                        lambda: -model.get_loss(logits_adv, tf.one_hot(maxclass, logits.shape[-1])),
                                        lambda:  model.get_loss(logits_adv, y_c[:n_adv])
                                        )
                loss_clean = model.get_loss(logits, y_c[n_adv:])
                reg_plain = regularizers.weight_decay()
                
                loss_ce = (loss_clean + loss_adv)/tf.cast(n_clean + n_adv, tf.float32)
                
                loss_tower = loss_ce + hps.lmbd * reg_plain

                update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)

                # Here are just different metrics for tensorboard
                loss_clean_metric, loss_clean_metric_upd = tf.metrics.mean(loss_clean/tf.cast(n_clean, tf.float32), name='loss_clean_metric')
                loss_out_metric, loss_out_metric_upd = tf.metrics.mean(loss_adv/tf.cast(n_adv, tf.float32), name='loss_out_metric')
                loss_total_metric, loss_total_metric_upd = tf.metrics.mean(loss_tower, name='loss_total_metric')
                reg_metric, reg_metric_upd = tf.metrics.mean(reg_plain, name='reg_metric')

                losses_clean_metric, losses_clean_metric_upds = losses_clean_metric + [loss_clean_metric], losses_clean_metric_upds + [loss_clean_metric_upd]
                losses_out_metric, losses_out_metric_upds = losses_out_metric + [loss_out_metric], losses_out_metric_upds + [loss_out_metric_upd]
                losses_total_metric, losses_total_metric_upds = losses_total_metric + [loss_total_metric], losses_total_metric_upds + [loss_total_metric_upd]
                regs_metric, regs_metric_upds = regs_metric + [reg_metric], regs_metric_upds + [reg_metric_upd]

                acc_rate, acc_rate_upd = tf.metrics.mean(tf.nn.in_top_k(
                    predictions=logits_c, targets=tf.argmax(y_c, 1), k=1),
                    name='top1_err_metric')
                avg_conf, avg_conf_upd = tf.metrics.mean(tf.reduce_max(probs_c, axis=1), name='avg_conf_metric')
                acc_rates, acc_rate_upds = acc_rates + [acc_rate], acc_rate_upds + [acc_rate_upd]
                avg_confs, avg_conf_upds = avg_confs + [avg_conf], avg_conf_upds + [avg_conf_upd]

                #Lists of the predictions
                logits_list += [logits_c]
                probs_list += [probs_c]
                labels_list += [y_c]           
                
    grads_vars = utils_tf.average_gradients(tower_grads)
    # Average error rates
    mean_loss_clean = tf.reduce_mean(tf.stack(losses_clean_metric))
    mean_loss_out = tf.reduce_mean(tf.stack(losses_out_metric))
    mean_loss_total = tf.reduce_mean(tf.stack(losses_total_metric))
    mean_reg = tf.reduce_mean(tf.stack(regs_metric))
    err_rate = 1 - tf.reduce_mean(tf.stack(acc_rates))
    avg_conf = tf.reduce_mean(tf.stack(avg_confs))
    mean_loss_clean_upd, mean_loss_out_upd, mean_loss_total_upd, mean_reg_upd = tf.group(*losses_clean_metric_upds), tf.group(*losses_out_metric_upds), tf.group(*losses_total_metric_upds), tf.group(*regs_metric_upds)
    err_rate_upd, avg_conf_upd = tf.group(*acc_rate_upds), tf.group(*avg_conf_upds)

    train_writer = tf.summary.FileWriter(tb_train, flush_secs=30)
    test_writer = tf.summary.FileWriter(tb_test, flush_secs=30)

    metric_vars = [var for var in tf.get_collection(tf.GraphKeys.LOCAL_VARIABLES) if 'metric' in var.op.name]
    init_metric_vars = tf.variables_initializer(var_list=metric_vars)
    saver = tf.train.Saver()

    tb_summaries_op = tf.summary.merge_all()

    gpu_options = tf.GPUOptions(visible_device_list=str(hps.gpus)[1:-1], per_process_gpu_memory_fraction=hps.gpu_memory)
    config = tf.ConfigProto(allow_soft_placement=True, gpu_options=gpu_options)