# Path: fully_supervised/lib/data.py
# DATASETS = create_datasets
# 
# Path: fully_supervised/lib/train.py
# class ClassifyFullySupervised(ClassifySemi):
#     """Fully supervised classification.
#     """
# 
#     def train_step(self, train_session, gen_labeled):
#         x = gen_labeled()
#         self.tmp.step = train_session.run([self.ops.train_op, self.ops.update_step],
#                                           feed_dict={self.ops.xt: 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.repeat().shuffle(FLAGS.shuffle).parse().augment()
#         train_labeled = train_labeled.batch(batch).prefetch(16).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()
#             gen_labeled = self.gen_labeled_fn(train_labeled)
#             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, gen_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))
# 
# Path: libml/models.py
# class CNN13(ClassifySemi):
# class ResNet(ClassifySemi):
# class ShakeNet(ClassifySemi):
# class MultiModel(CNN13, ResNet, ShakeNet):
#     def classifier(self, x, scales, filters, training, getter=None, **kwargs):
#     def classifier(self, x, scales, filters, repeat, training, getter=None, dropout=0, **kwargs):
#         def conv_args(k, f):
#         def residual(x0, filters, stride=1, activate_before_residual=False):
#     def classifier(self, x, scales, filters, repeat, training, getter=None, dropout=0, **kwargs):
#         def conv_args(k, f):
#         def residual(x0, filters, stride=1):
#             def branch():
#     def augment(self, x, l, smoothing, **kwargs):
#     def classifier(self, x, arch, **kwargs):
#     MODELS = ('cnn13', 'resnet', 'shake')
#     MODEL_CNN13, MODEL_RESNET, MODEL_SHAKE = MODELS
# 
# 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 combine_dicts(*args):
# 
# 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 fully_supervised.lib.data import DATASETS
from fully_supervised.lib.train import ClassifyFullySupervised
from libml import models, utils
from libml.utils import EasyDict
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Fully supervised training.
"""




FLAGS = flags.FLAGS


class FSBaseline(ClassifyFullySupervised, models.MultiModel):

    def model(self, batch, lr, wd, ema, **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')
        l_in = tf.placeholder(tf.int32, [batch], 'labels')
        wd *= lr
        classifier = lambda x, **kw: self.classifier(x, **kw, **kwargs).logits

        x, labels_x = self.augment(xt_in, tf.one_hot(l_in, self.nclass), **kwargs)
        logits_x = classifier(x, training=True)

        loss_xe = tf.nn.softmax_cross_entropy_with_logits_v2(labels=labels_x, logits=logits_x)
        loss_xe = tf.reduce_mean(loss_xe)
        tf.summary.scalar('losses/xe', loss_xe)

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

# Path: fully_supervised/lib/data.py
# DATASETS = create_datasets
# 
# Path: fully_supervised/lib/train.py
# class ClassifyFullySupervised(ClassifySemi):
#     """Fully supervised classification.
#     """
# 
#     def train_step(self, train_session, gen_labeled):
#         x = gen_labeled()
#         self.tmp.step = train_session.run([self.ops.train_op, self.ops.update_step],
#                                           feed_dict={self.ops.xt: 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.repeat().shuffle(FLAGS.shuffle).parse().augment()
#         train_labeled = train_labeled.batch(batch).prefetch(16).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()
#             gen_labeled = self.gen_labeled_fn(train_labeled)
#             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, gen_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))
# 
# Path: libml/models.py
# class CNN13(ClassifySemi):
# class ResNet(ClassifySemi):
# class ShakeNet(ClassifySemi):
# class MultiModel(CNN13, ResNet, ShakeNet):
#     def classifier(self, x, scales, filters, training, getter=None, **kwargs):
#     def classifier(self, x, scales, filters, repeat, training, getter=None, dropout=0, **kwargs):
#         def conv_args(k, f):
#         def residual(x0, filters, stride=1, activate_before_residual=False):
#     def classifier(self, x, scales, filters, repeat, training, getter=None, dropout=0, **kwargs):
#         def conv_args(k, f):
#         def residual(x0, filters, stride=1):
#             def branch():
#     def augment(self, x, l, smoothing, **kwargs):
#     def classifier(self, x, arch, **kwargs):
#     MODELS = ('cnn13', 'resnet', 'shake')
#     MODEL_CNN13, MODEL_RESNET, MODEL_SHAKE = MODELS
# 
# 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 combine_dicts(*args):
# 
# Path: libml/utils.py
# class EasyDict(dict):
#     def __init__(self, *args, **kwargs):
#         super(EasyDict, self).__init__(*args, **kwargs)
#         self.__dict__ = self


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

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Fully supervised training.
"""




FLAGS = flags.FLAGS


class FSBaseline(ClassifyFullySupervised, models.MultiModel):

    def model(self, batch, lr, wd, ema, **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')
        l_in = tf.placeholder(tf.int32, [batch], 'labels')
        wd *= lr
        classifier = lambda x, **kw: self.classifier(x, **kw, **kwargs).logits

        x, labels_x = self.augment(xt_in, tf.one_hot(l_in, self.nclass), **kwargs)
        logits_x = classifier(x, training=True)

        loss_xe = tf.nn.softmax_cross_entropy_with_logits_v2(labels=labels_x, logits=logits_x)
        loss_xe = tf.reduce_mean(loss_xe)
        tf.summary.scalar('losses/xe', loss_xe)

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

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