# Path: fully_supervised/lib/data.py
# DATASETS = {}
# 
# Path: fully_supervised/lib/train.py
# class ClassifyFullySupervised(ClassifySemi):
#     """Fully supervised classification.
#     """
# 
#     def train_step(self, train_session, data_labeled):
#         x = self.session.run(data_labeled)
#         self.tmp.step = train_session.run([self.ops.train_op, self.ops.update_step],
#                                           feed_dict={self.ops.x: x['image'],
#                                                      self.ops.label: x['label']})[1]
# 
#     def train(self, train_nimg, report_nimg):
#         if FLAGS.eval_ckpt:
#             self.eval_checkpoint(FLAGS.eval_ckpt)
#             return
#         batch = FLAGS.batch
#         train_labeled = self.dataset.train_labeled.batch(batch).prefetch(16)
#         train_labeled = train_labeled.make_one_shot_iterator().get_next()
#         scaffold = tf.train.Scaffold(saver=tf.train.Saver(max_to_keep=FLAGS.keep_ckpt,
#                                                           pad_step_number=10))
# 
#         with tf.Session(config=utils.get_config()) as sess:
#             self.session = sess
#             self.cache_eval()
# 
#         with tf.train.MonitoredTrainingSession(
#                 scaffold=scaffold,
#                 checkpoint_dir=self.checkpoint_dir,
#                 config=utils.get_config(),
#                 save_checkpoint_steps=FLAGS.save_kimg << 10,
#                 save_summaries_steps=report_nimg - batch) as train_session:
#             self.session = train_session._tf_sess()
#             self.tmp.step = self.session.run(self.step)
#             while self.tmp.step < train_nimg:
#                 loop = trange(self.tmp.step % report_nimg, report_nimg, batch,
#                               leave=False, unit='img', unit_scale=batch,
#                               desc='Epoch %d/%d' % (1 + (self.tmp.step // report_nimg), train_nimg // report_nimg))
#                 for _ in loop:
#                     self.train_step(train_session, train_labeled)
#                     while self.tmp.print_queue:
#                         loop.write(self.tmp.print_queue.pop(0))
#             while self.tmp.print_queue:
#                 print(self.tmp.print_queue.pop(0))
# 
#     def tune(self, train_nimg):
#         batch = FLAGS.batch
#         train_labeled = self.dataset.train_labeled.batch(batch).prefetch(16)
#         train_labeled = train_labeled.make_one_shot_iterator().get_next()
# 
#         for _ in trange(0, train_nimg, batch, leave=False, unit='img', unit_scale=batch, desc='Tuning'):
#             x = self.session.run([train_labeled])
#             self.session.run([self.ops.tune_op], feed_dict={self.ops.x: x['image'],
#                                                             self.ops.label: x['label']})
# 
# Path: libml/utils.py
# _GPUS = None
# FLAGS = flags.FLAGS
#         _GPUS = tuple([x.name for x in local_device_protos if x.device_type == 'GPU'])
# def get_config():
# 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 average_gradients(tower_grads):
# def para_list(fn, *args):
# def para_mean(fn, *args):
# def para_cat(fn, *args):
# 
# Path: libml/models.py
# class MultiModel(CNN13, ConvNet, ResNet, ShakeNet):
#     MODEL_CNN13, MODEL_CONVNET, MODEL_RESNET, MODEL_SHAKE = 'cnn13 convnet resnet shake'.split()
#     MODELS = MODEL_CNN13, MODEL_CONVNET, MODEL_RESNET, MODEL_SHAKE
# 
#     def augment(self, x, l, smoothing, **kwargs):
#         del kwargs
#         return x, l - smoothing * (l - 1. / self.nclass)
# 
#     def classifier(self, x, arch, **kwargs):
#         if arch == self.MODEL_CNN13:
#             return CNN13.classifier(self, x, **kwargs)
#         elif arch == self.MODEL_CONVNET:
#             return ConvNet.classifier(self, x, **kwargs)
#         elif arch == self.MODEL_RESNET:
#             return ResNet.classifier(self, x, **kwargs)
#         elif arch == self.MODEL_SHAKE:
#             return ShakeNet.classifier(self, x, **kwargs)
#         raise ValueError('Model %s does not exists, available ones are %s' % (arch, self.MODELS))


import functools
import os
import tensorflow as tf
from absl import app
from absl import flags
from easydict import EasyDict
from fully_supervised.lib.data import DATASETS
from fully_supervised.lib.train import ClassifyFullySupervised
from libml import utils
from libml.models import MultiModel
# See the License for the specific language governing permissions and
# limitations under the License.
"""Fully supervised training.
"""




FLAGS = flags.FLAGS


class FSBaseline(ClassifyFullySupervised, MultiModel):

    def model(self, lr, wd, ema, **kwargs):
        hwc = [self.dataset.height, self.dataset.width, self.dataset.colors]
        x_in = tf.placeholder(tf.float32, [None] + hwc, 'x')
        l_in = tf.placeholder(tf.int32, [None], 'labels')
        wd *= lr
        l = tf.one_hot(l_in, self.nclass)

        x, l = self.augment(x_in, l, **kwargs)
        classifier = functools.partial(self.classifier, **kwargs)
        logits = classifier(x, training=True)

        loss = tf.nn.softmax_cross_entropy_with_logits_v2(labels=l, logits=logits)
        loss = tf.reduce_mean(loss)
        tf.summary.scalar('losses/xe', loss)

        ema = tf.train.ExponentialMovingAverage(decay=ema)
        ema_op = ema.apply(utils.model_vars())

# Path: fully_supervised/lib/data.py
# DATASETS = {}
# 
# Path: fully_supervised/lib/train.py
# class ClassifyFullySupervised(ClassifySemi):
#     """Fully supervised classification.
#     """
# 
#     def train_step(self, train_session, data_labeled):
#         x = self.session.run(data_labeled)
#         self.tmp.step = train_session.run([self.ops.train_op, self.ops.update_step],
#                                           feed_dict={self.ops.x: x['image'],
#                                                      self.ops.label: x['label']})[1]
# 
#     def train(self, train_nimg, report_nimg):
#         if FLAGS.eval_ckpt:
#             self.eval_checkpoint(FLAGS.eval_ckpt)
#             return
#         batch = FLAGS.batch
#         train_labeled = self.dataset.train_labeled.batch(batch).prefetch(16)
#         train_labeled = train_labeled.make_one_shot_iterator().get_next()
#         scaffold = tf.train.Scaffold(saver=tf.train.Saver(max_to_keep=FLAGS.keep_ckpt,
#                                                           pad_step_number=10))
# 
#         with tf.Session(config=utils.get_config()) as sess:
#             self.session = sess
#             self.cache_eval()
# 
#         with tf.train.MonitoredTrainingSession(
#                 scaffold=scaffold,
#                 checkpoint_dir=self.checkpoint_dir,
#                 config=utils.get_config(),
#                 save_checkpoint_steps=FLAGS.save_kimg << 10,
#                 save_summaries_steps=report_nimg - batch) as train_session:
#             self.session = train_session._tf_sess()
#             self.tmp.step = self.session.run(self.step)
#             while self.tmp.step < train_nimg:
#                 loop = trange(self.tmp.step % report_nimg, report_nimg, batch,
#                               leave=False, unit='img', unit_scale=batch,
#                               desc='Epoch %d/%d' % (1 + (self.tmp.step // report_nimg), train_nimg // report_nimg))
#                 for _ in loop:
#                     self.train_step(train_session, train_labeled)
#                     while self.tmp.print_queue:
#                         loop.write(self.tmp.print_queue.pop(0))
#             while self.tmp.print_queue:
#                 print(self.tmp.print_queue.pop(0))
# 
#     def tune(self, train_nimg):
#         batch = FLAGS.batch
#         train_labeled = self.dataset.train_labeled.batch(batch).prefetch(16)
#         train_labeled = train_labeled.make_one_shot_iterator().get_next()
# 
#         for _ in trange(0, train_nimg, batch, leave=False, unit='img', unit_scale=batch, desc='Tuning'):
#             x = self.session.run([train_labeled])
#             self.session.run([self.ops.tune_op], feed_dict={self.ops.x: x['image'],
#                                                             self.ops.label: x['label']})
# 
# Path: libml/utils.py
# _GPUS = None
# FLAGS = flags.FLAGS
#         _GPUS = tuple([x.name for x in local_device_protos if x.device_type == 'GPU'])
# def get_config():
# 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 average_gradients(tower_grads):
# def para_list(fn, *args):
# def para_mean(fn, *args):
# def para_cat(fn, *args):
# 
# Path: libml/models.py
# class MultiModel(CNN13, ConvNet, ResNet, ShakeNet):
#     MODEL_CNN13, MODEL_CONVNET, MODEL_RESNET, MODEL_SHAKE = 'cnn13 convnet resnet shake'.split()
#     MODELS = MODEL_CNN13, MODEL_CONVNET, MODEL_RESNET, MODEL_SHAKE
# 
#     def augment(self, x, l, smoothing, **kwargs):
#         del kwargs
#         return x, l - smoothing * (l - 1. / self.nclass)
# 
#     def classifier(self, x, arch, **kwargs):
#         if arch == self.MODEL_CNN13:
#             return CNN13.classifier(self, x, **kwargs)
#         elif arch == self.MODEL_CONVNET:
#             return ConvNet.classifier(self, x, **kwargs)
#         elif arch == self.MODEL_RESNET:
#             return ResNet.classifier(self, x, **kwargs)
#         elif arch == self.MODEL_SHAKE:
#             return ShakeNet.classifier(self, x, **kwargs)
#         raise ValueError('Model %s does not exists, available ones are %s' % (arch, self.MODELS))


# Path: fully_supervised/fs_baseline.py
import functools
import os
import tensorflow as tf
from absl import app
from absl import flags
from easydict import EasyDict
from fully_supervised.lib.data import DATASETS
from fully_supervised.lib.train import ClassifyFullySupervised
from libml import utils
from libml.models import MultiModel

# See the License for the specific language governing permissions and
# limitations under the License.
"""Fully supervised training.
"""




FLAGS = flags.FLAGS


class FSBaseline(ClassifyFullySupervised, MultiModel):

    def model(self, lr, wd, ema, **kwargs):
        hwc = [self.dataset.height, self.dataset.width, self.dataset.colors]
        x_in = tf.placeholder(tf.float32, [None] + hwc, 'x')
        l_in = tf.placeholder(tf.int32, [None], 'labels')
        wd *= lr
        l = tf.one_hot(l_in, self.nclass)

        x, l = self.augment(x_in, l, **kwargs)
        classifier = functools.partial(self.classifier, **kwargs)
        logits = classifier(x, training=True)

        loss = tf.nn.softmax_cross_entropy_with_logits_v2(labels=l, logits=logits)
        loss = tf.reduce_mean(loss)
        tf.summary.scalar('losses/xe', loss)

        ema = tf.train.ExponentialMovingAverage(decay=ema)
        ema_op = ema.apply(utils.model_vars())

        ema_getter = functools.partial(utils.getter_ema, ema)