# Path: libml/utils.py
# _GPUS = None
# FLAGS = flags.FLAGS
#         _GPUS = tuple([x.name for x in local_device_protos if x.device_type == 'GPU'])
# class EasyDict(dict):
#     def __init__(self, *args, **kwargs):
# def get_config():
# def setup_main():
# def setup_tf():
# def smart_shape(x):
# def ilog2(x):
# def find_latest_checkpoint(dir, glob_term='model.ckpt-*.meta'):
# def get_latest_global_step(dir):
# def get_latest_global_step_in_subdir(dir):
# def getter_ema(ema, getter, name, *args, **kwargs):
# def model_vars(scope=None):
# def gpu(x):
# def get_available_gpus():
# def get_gpu():
# def average_gradients(tower_grads):
# def para_list(fn, *args):
# def para_mean(fn, *args):
# def para_cat(fn, *args):
# def interleave(x, batch):
# def de_interleave(x, batch):
# def combine_dicts(*args):
# 
# Path: libml/data.py
# DATA_DIR = None
# _DATA_CACHE = None
# SAMPLES_PER_CLASS = [1, 2, 3, 4, 5, 10, 25, 100, 400]
# FLAGS = flags.FLAGS
#     DATA_DIR = FLAGS.data_dir or os.environ['ML_DATA']
# DATASETS = functools.partial(create_datasets, augment_module.augment_function)
# PAIR_DATASETS = functools.partial(create_datasets, augment_module.pair_augment_function)
# MANY_DATASETS = functools.partial(create_datasets, augment_module.many_augment_function)
# QUAD_DATASETS = functools.partial(create_datasets, augment_module.quad_augment_function)
# def _data_setup():
# def record_parse_mnist(serialized_example, image_shape=None):
# def record_parse(serialized_example, image_shape=None):
# def compute_mean_std(data: tf.data.Dataset):
#         def iterator():
#     def __init__(self, data: tf.data.Dataset, augment_fn: AugmentPair, parse_fn=record_parse, image_shape=None):
#     def from_files(cls, filenames: list, augment_fn: AugmentPair, parse_fn=record_parse, image_shape=None):
#             def fetch_dataset(filename):
#     def empty_data(cls, image_shape, augment_fn: AugmentPair = None):
#         def _get_null_input(_):
#     def __getattr__(self, item):
#         def call_and_update(*args, **kwargs):
#     def parse(self):
#     def numpy_augment(self, *args, **kwargs):
#     def augment(self):
#     def memoize(self):
#         def tf_get(index, image_shape):
#             def get(index):
#     def __init__(self, name, train_labeled: DataSet, train_unlabeled: DataSet, test: DataSet, valid: DataSet,
#                  height=32, width=32, colors=3, nclass=10, mean=0, std=1, p_labeled=None, p_unlabeled=None):
#     def creator(cls, name, seed, label, valid, augment, parse_fn=record_parse, do_memoize=False,
#                 nclass=10, colors=3, height=32, width=32):
#         def create():
# def create_datasets(augment_fn):
# class DataSet:
# class DataSets:
# 
# Path: libml/utils.py
# class EasyDict(dict):
#     def __init__(self, *args, **kwargs):
#         super(EasyDict, self).__init__(*args, **kwargs)
#         self.__dict__ = self


import functools
import os
import tensorflow as tf
from absl import app
from absl import flags
from libml import utils, data, layers, models
from libml.utils import EasyDict
from third_party import vat_utils
class VAT(models.MultiModel):

    def model(self, batch, lr, wd, ema, warmup_pos, vat, vat_eps, entmin_weight, **kwargs):
        hwc = [self.dataset.height, self.dataset.width, self.dataset.colors]
        xt_in = tf.placeholder(tf.float32, [batch] + hwc, 'xt')  # For training
        x_in = tf.placeholder(tf.float32, [None] + hwc, 'x')
        y_in = tf.placeholder(tf.float32, [batch] + hwc, 'y')
        l_in = tf.placeholder(tf.int32, [batch], 'labels')
        wd *= lr
        warmup = tf.clip_by_value(tf.to_float(self.step) / (warmup_pos * (FLAGS.train_kimg << 10)), 0, 1)

        classifier = lambda x, **kw: self.classifier(x, **kw, **kwargs).logits
        l = tf.one_hot(l_in, self.nclass)
        logits_x = classifier(xt_in, training=True)
        post_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)  # Take only first call to update batch norm.
        logits_y = classifier(y_in, training=True)
        delta_y = vat_utils.generate_perturbation(y_in, logits_y, lambda x: classifier(x, training=True), vat_eps)
        logits_student = classifier(y_in + delta_y, training=True)
        logits_teacher = tf.stop_gradient(logits_y)
        loss_vat = layers.kl_divergence_from_logits(logits_student, logits_teacher)
        loss_vat = tf.reduce_mean(loss_vat)
        loss_entmin = tf.reduce_mean(tf.distributions.Categorical(logits=logits_y).entropy())

        loss = tf.nn.softmax_cross_entropy_with_logits_v2(labels=l, logits=logits_x)
        loss = tf.reduce_mean(loss)
        tf.summary.scalar('losses/xe', loss)
        tf.summary.scalar('losses/vat', loss_vat)
        tf.summary.scalar('losses/entmin', loss_entmin)

        ema = tf.train.ExponentialMovingAverage(decay=ema)

# Path: libml/utils.py
# _GPUS = None
# FLAGS = flags.FLAGS
#         _GPUS = tuple([x.name for x in local_device_protos if x.device_type == 'GPU'])
# class EasyDict(dict):
#     def __init__(self, *args, **kwargs):
# def get_config():
# def setup_main():
# def setup_tf():
# def smart_shape(x):
# def ilog2(x):
# def find_latest_checkpoint(dir, glob_term='model.ckpt-*.meta'):
# def get_latest_global_step(dir):
# def get_latest_global_step_in_subdir(dir):
# def getter_ema(ema, getter, name, *args, **kwargs):
# def model_vars(scope=None):
# def gpu(x):
# def get_available_gpus():
# def get_gpu():
# def average_gradients(tower_grads):
# def para_list(fn, *args):
# def para_mean(fn, *args):
# def para_cat(fn, *args):
# def interleave(x, batch):
# def de_interleave(x, batch):
# def combine_dicts(*args):
# 
# Path: libml/data.py
# DATA_DIR = None
# _DATA_CACHE = None
# SAMPLES_PER_CLASS = [1, 2, 3, 4, 5, 10, 25, 100, 400]
# FLAGS = flags.FLAGS
#     DATA_DIR = FLAGS.data_dir or os.environ['ML_DATA']
# DATASETS = functools.partial(create_datasets, augment_module.augment_function)
# PAIR_DATASETS = functools.partial(create_datasets, augment_module.pair_augment_function)
# MANY_DATASETS = functools.partial(create_datasets, augment_module.many_augment_function)
# QUAD_DATASETS = functools.partial(create_datasets, augment_module.quad_augment_function)
# def _data_setup():
# def record_parse_mnist(serialized_example, image_shape=None):
# def record_parse(serialized_example, image_shape=None):
# def compute_mean_std(data: tf.data.Dataset):
#         def iterator():
#     def __init__(self, data: tf.data.Dataset, augment_fn: AugmentPair, parse_fn=record_parse, image_shape=None):
#     def from_files(cls, filenames: list, augment_fn: AugmentPair, parse_fn=record_parse, image_shape=None):
#             def fetch_dataset(filename):
#     def empty_data(cls, image_shape, augment_fn: AugmentPair = None):
#         def _get_null_input(_):
#     def __getattr__(self, item):
#         def call_and_update(*args, **kwargs):
#     def parse(self):
#     def numpy_augment(self, *args, **kwargs):
#     def augment(self):
#     def memoize(self):
#         def tf_get(index, image_shape):
#             def get(index):
#     def __init__(self, name, train_labeled: DataSet, train_unlabeled: DataSet, test: DataSet, valid: DataSet,
#                  height=32, width=32, colors=3, nclass=10, mean=0, std=1, p_labeled=None, p_unlabeled=None):
#     def creator(cls, name, seed, label, valid, augment, parse_fn=record_parse, do_memoize=False,
#                 nclass=10, colors=3, height=32, width=32):
#         def create():
# def create_datasets(augment_fn):
# class DataSet:
# class DataSets:
# 
# Path: libml/utils.py
# class EasyDict(dict):
#     def __init__(self, *args, **kwargs):
#         super(EasyDict, self).__init__(*args, **kwargs)
#         self.__dict__ = self


# Path: vat.py
import functools
import os
import tensorflow as tf
from absl import app
from absl import flags
from libml import utils, data, layers, models
from libml.utils import EasyDict
from third_party import vat_utils

class VAT(models.MultiModel):

    def model(self, batch, lr, wd, ema, warmup_pos, vat, vat_eps, entmin_weight, **kwargs):
        hwc = [self.dataset.height, self.dataset.width, self.dataset.colors]
        xt_in = tf.placeholder(tf.float32, [batch] + hwc, 'xt')  # For training
        x_in = tf.placeholder(tf.float32, [None] + hwc, 'x')
        y_in = tf.placeholder(tf.float32, [batch] + hwc, 'y')
        l_in = tf.placeholder(tf.int32, [batch], 'labels')
        wd *= lr
        warmup = tf.clip_by_value(tf.to_float(self.step) / (warmup_pos * (FLAGS.train_kimg << 10)), 0, 1)

        classifier = lambda x, **kw: self.classifier(x, **kw, **kwargs).logits
        l = tf.one_hot(l_in, self.nclass)
        logits_x = classifier(xt_in, training=True)
        post_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)  # Take only first call to update batch norm.
        logits_y = classifier(y_in, training=True)
        delta_y = vat_utils.generate_perturbation(y_in, logits_y, lambda x: classifier(x, training=True), vat_eps)
        logits_student = classifier(y_in + delta_y, training=True)
        logits_teacher = tf.stop_gradient(logits_y)
        loss_vat = layers.kl_divergence_from_logits(logits_student, logits_teacher)
        loss_vat = tf.reduce_mean(loss_vat)
        loss_entmin = tf.reduce_mean(tf.distributions.Categorical(logits=logits_y).entropy())

        loss = tf.nn.softmax_cross_entropy_with_logits_v2(labels=l, logits=logits_x)
        loss = tf.reduce_mean(loss)
        tf.summary.scalar('losses/xe', loss)
        tf.summary.scalar('losses/vat', loss_vat)
        tf.summary.scalar('losses/entmin', loss_entmin)

        ema = tf.train.ExponentialMovingAverage(decay=ema)

        ema_op = ema.apply(utils.model_vars())