def main(args):
    global best_acc
    best_acc = 0
    with tf.Graph().as_default():
        print('Input data preprocessing...')
        with tf.name_scope('DataPreprocess'):
            r_train = (500.0 / 6.0)
            r_test = (100.0 / 6.0)
            nClass = 100
            mon_instance = 2498.0
            unClass = 0
            unmon_instance = unClass
            dim = 5000
            with tf.device('/cpu:0'):
                (train_x, train_y, test_x_data, test_y_data) = split_awf_closed(r_train, r_test, nClass, mon_instance, unmon_instance, dim)

            def reshape_and_scale(x, img_shape=((- 1), dim, 1)):
                return x.reshape(img_shape).astype(np.float32)
            train_x = reshape_and_scale(train_x)
            test_x_data = reshape_and_scale(test_x_data)
            awf_data2 = np.load((HOME + '/datasets/awf2.npz'), allow_pickle=True)
            train_x_unlabeled = awf_data2['data']
            train_y_unlabeled = awf_data2['labels']
            train_x_unlabeled = reshape_and_scale(train_x_unlabeled)
            (X, y) = shuffle(train_x, train_y)
            print(X.shape)
            print(y.shape)
        print('Setup the input pipeline...')
        with tf.name_scope('InputPipeline'):
            (train_x_labeled, train_y_labeled) = ([], [])
            for i in range(args.num_classes):
                print(i)
                train_x_labeled.append(X[(y == i)][:args.num_labeled_examples])
                train_y_labeled.append(y[(y == i)][:args.num_labeled_examples])
            train_x_labeled_data = np.concatenate(train_x_labeled)
            train_y_labeled_data = np.concatenate(train_y_labeled)
            train_x_unlabeled_data = train_x_unlabeled
            train_y_unlabeled_data = train_y_unlabeled
            (train_x_unlabeled2, train_y_unlabeled2) = shuffle(train_x_unlabeled, train_y_unlabeled)
            train_x_unlabeled2_data = train_x_unlabeled2
            train_y_unlabeled2_data = train_y_unlabeled2
            labeled_X = tf.placeholder(tf.float32, shape=[None, dim, 1])
            labeled_y = tf.placeholder(tf.int64, shape=[None])
            unlabeled_X = tf.placeholder(tf.float32, shape=[None, dim, 1])
            unlabeled_y = tf.placeholder(tf.int64, shape=[None])
            unlabeled_X2 = tf.placeholder(tf.float32, shape=[None, dim, 1])
            unlabeled_y2 = tf.placeholder(tf.int64, shape=[None])
            test_X = tf.placeholder(tf.float32, shape=[None, dim, 1])
            test_y = tf.placeholder(tf.int64, shape=[None])
            train_labeled_dataset = tf.data.Dataset.from_tensor_slices((labeled_X, labeled_y)).shuffle(buffer_size=len(train_x_labeled_data)).repeat()
            train_labeled_dataset = train_labeled_dataset.batch(args.batch_size)
            iterator_labeled = train_labeled_dataset.make_initializable_iterator()
            (traces_lab, labels_lab) = iterator_labeled.get_next()
            train_unlabeled_dataset = tf.data.Dataset.from_tensor_slices((unlabeled_X, unlabeled_y, unlabeled_X2, unlabeled_y2)).shuffle(buffer_size=len(train_x_labeled_data)).repeat()
            train_unlabeled_dataset = train_unlabeled_dataset.batch(args.batch_size)
            iterator_unlabeled = train_unlabeled_dataset.make_initializable_iterator()
            (traces_unl, labels_unl, traces_unl2, labels_unl2) = iterator_unlabeled.get_next()
            test_dataset = tf.data.Dataset.from_tensor_slices((test_X, test_y)).repeat()
            test_dataset = test_dataset.batch(args.batch_size)
            iterator_test = test_dataset.make_initializable_iterator()
            (traces_test, labels_test) = iterator_test.get_next()
        with tf.name_scope('BatchSize'):
            batch_size_tensor = tf.shape(traces_lab)[0]
        (z, z_perturbed) = define_noise(batch_size_tensor, args)
        with tf.name_scope('Generator'):
            g_model = define_generator()
            traces_fake = g_model(z)
            traces_fake_perturbed = g_model(z_perturbed)
        with tf.name_scope('Discriminator') as discriminator_scope:
            d_model = Discriminator()
            (logits_fake, features_fake) = d_model(traces_fake, training=True)
            (logits_fake_perturbed, _) = d_model(traces_fake_perturbed, training=True)
            (logits_real_unl, features_real_unl) = d_model(traces_unl, training=True)
            (logits_real_lab, features_real_lab) = d_model(traces_lab, training=True)
            (logits_train, _) = d_model(traces_lab, training=False)
        with tf.name_scope('DiscriminatorLoss'):
            loss_supervised = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels_lab, logits=logits_real_lab))
            logits_sum_real = tf.reduce_logsumexp(logits_real_unl, axis=1)
            logits_sum_fake = tf.reduce_logsumexp(logits_fake, axis=1)
            loss_unsupervised = (0.5 * ((tf.negative(tf.reduce_mean(logits_sum_real)) + tf.reduce_mean(tf.nn.softplus(logits_sum_real))) + tf.reduce_mean(tf.nn.softplus(logits_sum_fake))))
            loss_d = (loss_supervised + loss_unsupervised)
            if args.man_reg:
                loss_d += ((0.001 * tf.nn.l2_loss((logits_fake - logits_fake_perturbed))) / tf.to_float(batch_size_tensor))
        with tf.name_scope('Train') as train_scope:
            optimizer = tf.train.AdamOptimizer((args.lr * 0.25))
            optimize_d = optimizer.minimize(loss_d, var_list=d_model.trainable_variables)
            train_accuracy_op = accuracy(logits_train, labels_lab)
        with tf.name_scope(discriminator_scope):
            with tf.control_dependencies([optimize_d]):
                (logits_fake, features_fake) = d_model(traces_fake, training=True)
                (logits_real_unl, features_real_unl) = d_model(traces_unl2, training=True)
        with tf.name_scope('GeneratorLoss'):
            feature_mean_real = tf.reduce_mean(features_real_unl, axis=0)
            feature_mean_fake = tf.reduce_mean(features_fake, axis=0)
            loss_g = tf.reduce_mean(tf.abs((feature_mean_real - feature_mean_fake)))
        with tf.name_scope(train_scope):
            optimizer = tf.train.AdamOptimizer(args.lr, beta1=0.5)
            train_op = optimizer.minimize(loss_g, var_list=g_model.trainable_variables)
        with tf.name_scope(discriminator_scope):
            with tf.name_scope('Test'):
                (logits_test, _) = d_model(traces_test, training=False)
                test_accuracy_op = accuracy(logits_test, labels_test)
        with tf.name_scope('Summaries'):
            summary_op = tf.summary.merge([tf.summary.scalar('LossDiscriminator', loss_d), tf.summary.scalar('LossGenerator', loss_g), tf.summary.scalar('ClassificationAccuracyTrain', train_accuracy_op), tf.summary.scalar('ClassificationAccuracyTest', test_accuracy_op)])
        writer = tf.summary.FileWriter(_next_logdir('tensorboard/wfi5_cw'))
        print('Run training...')
        steps_per_epoch = ((len(train_x_labeled_data) + len(train_x_unlabeled_data)) // args.batch_size)
        steps_per_test = (test_x_data.shape[0] // args.batch_size)
        with tf.Session() as sess:
            sess.run(tf.global_variables_initializer())
            step = 0
            for epoch in range(args.train_epochs):
                (losses_d, losses_g, accuracies) = ([], [], [])
                print('Epoch {}'.format(epoch))
                pbar = tqdm.trange(steps_per_epoch)
                sess.run(iterator_labeled.initializer, feed_dict={labeled_X: train_x_labeled_data, labeled_y: train_y_labeled_data})
                sess.run(iterator_unlabeled.initializer, feed_dict={unlabeled_X: train_x_unlabeled_data, unlabeled_y: train_y_unlabeled_data, unlabeled_X2: train_x_unlabeled2_data, unlabeled_y2: train_y_unlabeled2_data})
                sess.run(iterator_test.initializer, feed_dict={test_X: test_x_data, test_y: test_y_data})
                for _ in pbar:
                    if ((step % 1000) == 0):
                        (_, loss_g_batch, loss_d_batch, summ, accuracy_batch) = sess.run([train_op, loss_g, loss_d, summary_op, train_accuracy_op])
                        writer.add_summary(summ, global_step=step)
                    else:
                        (_, loss_g_batch, loss_d_batch, accuracy_batch) = sess.run([train_op, loss_g, loss_d, train_accuracy_op])
                    pbar.set_description('Discriminator loss {0:.3f}, Generator loss {1:.3f}'.format(loss_d_batch, loss_g_batch))
                    losses_d.append(loss_d_batch)
                    losses_g.append(loss_g_batch)
                    accuracies.append(accuracy_batch)
                    step += 1
                print('Discriminator loss: {0:.4f}, Generator loss: {1:.4f}, Train accuracy: {2:.4f}'.format(np.mean(losses_d), np.mean(losses_g), np.mean(accuracies)))
                accuracies = [sess.run(test_accuracy_op) for _ in range(steps_per_test)]
                if (np.mean(accuracies) > best_acc):
                    best_acc = np.mean(accuracies)
                if (epoch == (int(args.train_epochs) - 1)):
                    print('Test accuracy: {0:.4f}'.format(np.mean(accuracies)))
                    print('Best accuracy: {0:.4f}'.format(best_acc))
def main(args):
    global best_acc
    best_acc = 0
    with tf.Graph().as_default():
        print('Input data preprocessing...')
        with tf.name_scope('DataPreprocess'):
            r_train = (500.0 / 6.0)
            r_test = (100.0 / 6.0)
            nClass = 100
            mon_instance = 2498.0
            unClass = 0
            unmon_instance = unClass
            dim = 5000
            with tf.device('/cpu:0'):
                (train_x, train_y, test_x_data, test_y_data) = split_awf_closed(r_train, r_test, nClass, mon_instance, unmon_instance, dim)

            def reshape_and_scale(x, img_shape=((- 1), dim, 1)):
                return x.reshape(img_shape).astype(np.float32)
            train_x = reshape_and_scale(train_x)
            test_x_data = reshape_and_scale(test_x_data)
            awf_data2 = np.load((HOME + '/datasets/awf2.npz'), allow_pickle=True)
            train_x_unlabeled = awf_data2['data']
            train_y_unlabeled = awf_data2['labels']
            train_x_unlabeled = reshape_and_scale(train_x_unlabeled)
            (X, y) = shuffle(train_x, train_y)
            print(X.shape)
            print(y.shape)
        print('Setup the input pipeline...')
        with tf.name_scope('InputPipeline'):
            (train_x_labeled, train_y_labeled) = ([], [])
            for i in range(args.num_classes):
                print(i)
                train_x_labeled.append(X[(y == i)][:args.num_labeled_examples])
                train_y_labeled.append(y[(y == i)][:args.num_labeled_examples])
            train_x_labeled_data = np.concatenate(train_x_labeled)
            train_y_labeled_data = np.concatenate(train_y_labeled)
            train_x_unlabeled_data = train_x_unlabeled
            train_y_unlabeled_data = train_y_unlabeled
            (train_x_unlabeled2, train_y_unlabeled2) = shuffle(train_x_unlabeled, train_y_unlabeled)
            train_x_unlabeled2_data = train_x_unlabeled2
            train_y_unlabeled2_data = train_y_unlabeled2
            labeled_X = tf.placeholder(tf.float32, shape=[None, dim, 1])
            labeled_y = tf.placeholder(tf.int64, shape=[None])
            unlabeled_X = tf.placeholder(tf.float32, shape=[None, dim, 1])
            unlabeled_y = tf.placeholder(tf.int64, shape=[None])
            unlabeled_X2 = tf.placeholder(tf.float32, shape=[None, dim, 1])
            unlabeled_y2 = tf.placeholder(tf.int64, shape=[None])
            test_X = tf.placeholder(tf.float32, shape=[None, dim, 1])
            test_y = tf.placeholder(tf.int64, shape=[None])
            train_labeled_dataset = tf.data.Dataset.from_tensor_slices((labeled_X, labeled_y)).shuffle(buffer_size=len(train_x_labeled_data)).repeat()
            train_labeled_dataset = train_labeled_dataset.batch(args.batch_size)
            iterator_labeled = train_labeled_dataset.make_initializable_iterator()
            (traces_lab, labels_lab) = iterator_labeled.get_next()
            train_unlabeled_dataset = tf.data.Dataset.from_tensor_slices((unlabeled_X, unlabeled_y, unlabeled_X2, unlabeled_y2)).shuffle(buffer_size=len(train_x_labeled_data)).repeat()
            train_unlabeled_dataset = train_unlabeled_dataset.batch(args.batch_size)
            iterator_unlabeled = train_unlabeled_dataset.make_initializable_iterator()
            (traces_unl, labels_unl, traces_unl2, labels_unl2) = iterator_unlabeled.get_next()
            test_dataset = tf.data.Dataset.from_tensor_slices((test_X, test_y)).repeat()
            test_dataset = test_dataset.batch(args.batch_size)
            iterator_test = test_dataset.make_initializable_iterator()
            (traces_test, labels_test) = iterator_test.get_next()
        with tf.name_scope('BatchSize'):
            batch_size_tensor = tf.shape(traces_lab)[0]
        (z, z_perturbed) = define_noise(batch_size_tensor, args)
        with tf.name_scope('Generator'):
            g_model = define_generator()
            traces_fake = g_model(z)
            traces_fake_perturbed = g_model(z_perturbed)
        with tf.name_scope('Discriminator') as discriminator_scope:
            d_model = Discriminator()
            (logits_fake, features_fake) = d_model(traces_fake, training=True)
            (logits_fake_perturbed, _) = d_model(traces_fake_perturbed, training=True)
            (logits_real_unl, features_real_unl) = d_model(traces_unl, training=True)
            (logits_real_lab, features_real_lab) = d_model(traces_lab, training=True)
            (logits_train, _) = d_model(traces_lab, training=False)
        with tf.name_scope('DiscriminatorLoss'):
            loss_supervised = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels_lab, logits=logits_real_lab))
            logits_sum_real = tf.reduce_logsumexp(logits_real_unl, axis=1)
            logits_sum_fake = tf.reduce_logsumexp(logits_fake, axis=1)
            loss_unsupervised = (0.5 * ((tf.negative(tf.reduce_mean(logits_sum_real)) + tf.reduce_mean(tf.nn.softplus(logits_sum_real))) + tf.reduce_mean(tf.nn.softplus(logits_sum_fake))))
            loss_d = (loss_supervised + loss_unsupervised)
            if args.man_reg:
                loss_d += ((0.001 * tf.nn.l2_loss((logits_fake - logits_fake_perturbed))) / tf.to_float(batch_size_tensor))
        with tf.name_scope('Train') as train_scope:
            optimizer = tf.train.AdamOptimizer((args.lr * 0.25))
            optimize_d = optimizer.minimize(loss_d, var_list=d_model.trainable_variables)
            train_accuracy_op = accuracy(logits_train, labels_lab)
        with tf.name_scope(discriminator_scope):
            with tf.control_dependencies([optimize_d]):
                (logits_fake, features_fake) = d_model(traces_fake, training=True)
                (logits_real_unl, features_real_unl) = d_model(traces_unl2, training=True)
        with tf.name_scope('GeneratorLoss'):
            feature_mean_real = tf.reduce_mean(features_real_unl, axis=0)
            feature_mean_fake = tf.reduce_mean(features_fake, axis=0)
            loss_g = tf.reduce_mean(tf.abs((feature_mean_real - feature_mean_fake)))
        with tf.name_scope(train_scope):
            optimizer = tf.train.AdamOptimizer(args.lr, beta1=0.5)
            train_op = optimizer.minimize(loss_g, var_list=g_model.trainable_variables)
        with tf.name_scope(discriminator_scope):
            with tf.name_scope('Test'):
                (logits_test, _) = d_model(traces_test, training=False)
                test_accuracy_op = accuracy(logits_test, labels_test)
        with tf.name_scope('Summaries'):
            summary_op = tf.summary.merge([tf.summary.scalar('LossDiscriminator', loss_d), tf.summary.scalar('LossGenerator', loss_g), tf.summary.scalar('ClassificationAccuracyTrain', train_accuracy_op), tf.summary.scalar('ClassificationAccuracyTest', test_accuracy_op)])
        writer = tf.summary.FileWriter(_next_logdir('tensorboard/wfi5_cw'))
        print('Run training...')
        steps_per_epoch = ((len(train_x_labeled_data) + len(train_x_unlabeled_data)) // args.batch_size)
        steps_per_test = (test_x_data.shape[0] // args.batch_size)
        with tf.Session() as sess:
            sess.run(tf.global_variables_initializer())
            step = 0
            for epoch in range(args.train_epochs):
                (losses_d, losses_g, accuracies) = ([], [], [])
                print('Epoch {}'.format(epoch))
                pbar = tqdm.trange(steps_per_epoch)
                sess.run(iterator_labeled.initializer, feed_dict={labeled_X: train_x_labeled_data, labeled_y: train_y_labeled_data})
                sess.run(iterator_unlabeled.initializer, feed_dict={unlabeled_X: train_x_unlabeled_data, unlabeled_y: train_y_unlabeled_data, unlabeled_X2: train_x_unlabeled2_data, unlabeled_y2: train_y_unlabeled2_data})
                sess.run(iterator_test.initializer, feed_dict={test_X: test_x_data, test_y: test_y_data})
                for _ in pbar:
                    if ((step % 1000) == 0):
                        (_, loss_g_batch, loss_d_batch, summ, accuracy_batch) = sess.run([train_op, loss_g, loss_d, summary_op, train_accuracy_op])
                        writer.add_summary(summ, global_step=step)
                    else:
                        (_, loss_g_batch, loss_d_batch, accuracy_batch) = sess.run([train_op, loss_g, loss_d, train_accuracy_op])
                    pbar.set_description('Discriminator loss {0:.3f}, Generator loss {1:.3f}'.format(loss_d_batch, loss_g_batch))
                    losses_d.append(loss_d_batch)
                    losses_g.append(loss_g_batch)
                    accuracies.append(accuracy_batch)
                    step += 1
                print('Discriminator loss: {0:.4f}, Generator loss: {1:.4f}, Train accuracy: {2:.4f}'.format(np.mean(losses_d), np.mean(losses_g), np.mean(accuracies)))
                accuracies = [sess.run(test_accuracy_op) for _ in range(steps_per_test)]
                if (np.mean(accuracies) > best_acc):
                    best_acc = np.mean(accuracies)
                if (epoch == (int(args.train_epochs) - 1)):
                    print('Test accuracy: {0:.4f}'.format(np.mean(accuracies)))
                    print('Best accuracy: {0:.4f}'.format(best_acc))