text:
# coding=utf-8
# Copyright 2019 The Meta-Dataset Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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.

# Lint as: python2, python3
"""Interface for a learner that uses BenchmarkReaderDataSource to get data."""
# TODO(lamblinp): Update variable names to be more consistent
# - target, class_idx, label
# - support, query

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import os

from absl import logging
import gin.tf
from meta_dataset import learner
from meta_dataset.data import dataset_spec as dataset_spec_lib
from meta_dataset.data import learning_spec
from meta_dataset.data import pipeline
from meta_dataset.data import providers
import numpy as np
import six
from six.moves import range
from six.moves import zip
import tensorflow as tf

from tensorflow.core.protobuf import rewriter_config_pb2  # pylint: disable=g-direct-tensorflow-import

# Enable TensorFlow optimizations. It can add a few minutes to the first
# calls to session.run(), but decrease memory usage.
ENABLE_TF_OPTIMIZATIONS = True
# Enable tf.data optimizations, which are applied to the input data pipeline.
# It may be helpful to disable them when investigating regressions due to
# changes in tf.data (see b/121130181 for instance), but they seem to be helpful
# (or at least not detrimental) in general.
ENABLE_DATA_OPTIMIZATIONS = True

EMBEDDING_KEYWORDS = ('conv', 'resnet')

DATASETS_WITH_EXAMPLE_SPLITS = ()
TF_DATA_OPTIONS = tf.data.Options()
if not ENABLE_DATA_OPTIMIZATIONS:
  # The Options object can be used to control which static or dynamic
  # optimizations to apply.
  TF_DATA_OPTIONS.experimental_optimization.apply_default_optimizations = False

NAME_TO_LEARNER = {
    'Baseline': learner.BaselineLearner,
    'BaselineFinetune': learner.BaselineFinetuneLearner,
    'MatchingNet': learner.MatchingNetworkLearner,
    'PrototypicalNet': learner.PrototypicalNetworkLearner,
    'MAML': learner.MAMLLearner,
    'RelationNet': learner.RelationNetworkLearner,
}
BATCH_LEARNER_NAMES = ['Baseline', 'BaselineFinetune']
EPISODIC_LEARNER_NAMES = [
    'MatchingNet', 'PrototypicalNet', 'MAML', 'RelationNet'
]
BATCH_LEARNERS = [NAME_TO_LEARNER[name] for name in BATCH_LEARNER_NAMES]
EPISODIC_LEARNERS = [NAME_TO_LEARNER[name] for name in EPISODIC_LEARNER_NAMES]

# TODO(eringrant): Use `learning_spec.Split.TRAIN`, `learning_spec.Split.VALID`,
# and `learning_spec.Split.TEST` instead of string constants, and replace all
# remaining string redefinitions.
TRAIN_SPLIT = 'train'
VALID_SPLIT = 'valid'
TEST_SPLIT = 'test'

class UnexpectedSplitError(ValueError):

  def __init__(self,
               unexpected_split,
               expected_splits=(TRAIN_SPLIT, TEST_SPLIT, VALID_SPLIT)):
    super(UnexpectedSplitError,
          self).__init__('Split must be one of {}, but received `{}`. '.format(
              expected_splits, unexpected_split))

@gin.configurable('benchmark')
def get_datasets_and_restrictions(train_datasets='',
                                  eval_datasets='',
                                  restrict_classes=None,
                                  restrict_num_per_class=None):
  """Gets the list of dataset names and possible restrictions on their classes.

  Args:
    train_datasets: A string of comma-separated dataset names for training.
    eval_datasets: A string of comma-separated dataset names for evaluation.
    restrict_classes: If provided, a dict that maps dataset names to a dict that
      specifies for each of `meta_dataset.trainer.TRAIN_SPLIT`,
      `meta_dataset.trainer.VALID_SPLIT` and `meta_dataset.trainer.TEST_SPLIT`
      the number of classes to restrict to. This can lead to some classes of a
      particular split of a particular dataset never participating in episode
      creation.
    restrict_num_per_class: If provided, a dict that maps dataset names to a
      dict that specifies for each of `meta_dataset.trainer.TRAIN_SPLIT`,
      `meta_dataset.trainer.VALID_SPLIT` and `meta_dataset.trainer.TEST_SPLIT`
      the number of examples per class to restrict to. For datasets / splits
      that are not specified, no restriction is applied.

  Returns:
    Two lists of dataset names and two possibly empty dictionaries.
  """
  if restrict_classes is None:
    restrict_classes = {}
  if restrict_num_per_class is None:
    restrict_num_per_class = {}

  train_datasets = [d.strip() for d in train_datasets.split(',')]
  eval_datasets = [d.strip() for d in eval_datasets.split(',')]

  return train_datasets, eval_datasets, restrict_classes, restrict_num_per_class

def apply_dataset_options(dataset):
  """Apply the module-wide set of dataset options to dataset.

  In particular, this is used to enable or disable tf.data optimizations.
  This applies to the whole pipeline, so we can just set it at the end.

  Args:
    dataset: a tf.data.Dataset object.

  Returns:
    A tf.data.Dataset object with options applied.
  """
  return dataset.with_options(TF_DATA_OPTIONS)

def compute_train_class_proportions(episode, shots, dataset_spec):
  """Computes the proportion of each class' examples in the support set.

  Args:
    episode: An EpisodeDataset.
    shots: A 1D Tensor whose length is the `way' of the episode that stores the
      shots for this episode.
    dataset_spec: A DatasetSpecification.

  Returns:
    class_props: A 1D Tensor whose length is the `way' of the episode, storing
      for each class the proportion of its examples that are in the support set.
  """
  # Get the total number of examples of each class in the dataset.
  num_dataset_classes = len(dataset_spec.images_per_class)
  num_images_per_class = [
      dataset_spec.get_total_images_per_class(class_id)
      for class_id in range(num_dataset_classes)
  ]

  # Get the (absolute) class ID's that appear in the episode.
  class_ids, _ = tf.unique(episode.train_class_ids)  # [?, ]

  # Make sure that class_ids are valid indices of num_images_per_class. This is
  # important since tf.gather will fail silently and return zeros otherwise.
  num_classes = tf.shape(num_images_per_class)[0]
  check_valid_inds_op = tf.assert_less(class_ids, num_classes)
  with tf.control_dependencies([check_valid_inds_op]):
    # Get the total number of examples of each class that is in the episode.
    num_images_per_class = tf.gather(num_images_per_class, class_ids)  # [?, ]

  # Get the proportions of examples of each class that appear in the train set.
  class_props = tf.truediv(shots, num_images_per_class)
  return class_props

def get_split_enum(split):
  """Returns the Enum value corresponding to the given split.

  Args:
    split: A String.

  Raises:
    UnexpectedSplitError: split not TRAIN_SPLIT, VALID_SPLIT, or TEST_SPLIT.
  """
  # Get the int representing the chosen split.
  if split == TRAIN_SPLIT:
    split_enum = learning_spec.Split.TRAIN
  elif split == VALID_SPLIT:
    split_enum = learning_spec.Split.VALID
  elif split == TEST_SPLIT:
    split_enum = learning_spec.Split.TEST
  else:
    raise UnexpectedSplitError(split)
  return split_enum

def compute_episode_stats(episode):
  """Computes various episode stats: way, shots, and class IDs.

  Args:
    episode: An EpisodeDataset.

  Returns:
    way: An int constant tensor. The number of classes in the episode.
    shots: An int 1D tensor: The number of support examples per class.
    class_ids: An int 1D tensor: (absolute) class IDs.
  """
  # The train labels of the next episode.
  train_labels = episode.train_labels
  # Compute way.
  episode_classes, _ = tf.unique(train_labels)
  way = tf.size(episode_classes)
  # Compute shots.
  class_ids = tf.reshape(tf.range(way), [way, 1])
  class_labels = tf.reshape(train_labels, [1, -1])
  is_equal = tf.equal(class_labels, class_ids)
  shots = tf.reduce_sum(tf.cast(is_equal, tf.int32), axis=1)
  # Compute class_ids.
  class_ids, _ = tf.unique(episode.train_class_ids)
  return way, shots, class_ids

@gin.configurable
class LearnConfig(object):
  """A collection of values pertaining to learning."""

  def __init__(self, num_updates, batch_size, num_eval_episodes,
               checkpoint_every, validate_every, log_every,
               transductive_batch_norm):
    """Initializes a LearnConfig.

    Args:
      num_updates: An integer, the number of training updates.
      batch_size: An integer, the size of batches for non-episodic models.
      num_eval_episodes: An integer, the number of episodes for evaluation.
      checkpoint_every: An integer, the number of episodes between consecutive
        checkpoints.
      validate_every: An integer, the number of episodes between consecutive
        validatations.
      log_every: An integer, the number of episodes between consecutive logging.
      transductive_batch_norm: Whether to batch normalize in the transductive
        setting where the mean and variance for normalization are computed from
        both the support and query sets.
    """
    self.num_updates = num_updates
    self.batch_size = batch_size
    self.num_eval_episodes = num_eval_episodes
    self.checkpoint_every = checkpoint_every
    self.validate_every = validate_every
    self.log_every = log_every
    self.transductive_batch_norm = transductive_batch_norm

# TODO(manzagop): consider moving here num_updates/batch_size from LearnConfig.
# TODO(manzagop): consider combining train and eval learner and inferring
#   episodic.
@gin.configurable
class LearnerConfig(object):
  """A collection of values pertaining to the model."""

  def __init__(
      self,
      episodic,
      train_learner,
      eval_learner,
      pretrained_checkpoint,
      checkpoint_for_eval,
      embedding_network,
      learning_rate,
      decay_learning_rate,
      decay_every,
      decay_rate,
      experiment_name,
      pretrained_source,
  ):
    # pyformat: disable
    """Initializes a LearnerConfig.

    Args:
      episodic: A boolean, whether meta-training is episodic.
      train_learner: A string, the name of the learner to use for meta-training.
      eval_learner: A string, the name of the learner to use for
        meta-evaluation.
      pretrained_checkpoint: A string, the path to a checkpoint to use for
        initializing a model prior to training.
      checkpoint_for_eval: A string, the path to a checkpoint to restore for
        evaluation.
      embedding_network: A string, the embedding network to use.
      learning_rate: A float, the meta-learning learning rate.
      decay_learning_rate: A boolean, whether to decay the learning rate.
      decay_every: An integer, the learning rate is decayed for every multiple
        of this value.
      decay_rate: A float, the decay to apply to the learning rate.
      experiment_name: A string, a name for the experiment.
      pretrained_source: A string, the pretraining setup to use.
    """
    # pyformat: enable
    if checkpoint_for_eval and pretrained_checkpoint:
      raise ValueError(
          'Cannot define both `checkpoint_for_eval` and '
          '`pretrained_checkpoint`. The difference between them is that for '
          'the former, all variables are restored (including the global step), '
          'while the latter is only applicable to the start of training for '
          'initializing the model from pre-trained weights. It is also only '
          'applicable to episodic models and restores only the embedding '
          'weights.')

    self.episodic = episodic
    self.train_learner = train_learner
    self.eval_learner = eval_learner
    self.pretrained_checkpoint = pretrained_checkpoint
    self.checkpoint_for_eval = checkpoint_for_eval
    self.embedding_network = embedding_network
    self.learning_rate = learning_rate
    self.decay_learning_rate = decay_learning_rate
    self.decay_every = decay_every
    self.decay_rate = decay_rate
    self.experiment_name = experiment_name
    self.pretrained_source = pretrained_source

@gin.configurable
class Trainer(object):
  """A Trainer for training a Learner on data provided by ReaderDataSource."""

  def __init__(
      self,
      train_learner,
      eval_learner,
      is_training,
      train_dataset_list,
      eval_dataset_list,
      restrict_classes,
      restrict_num_per_class,
      checkpoint_dir,
      summary_dir,
      records_root_dir,
      eval_finegrainedness,
      eval_finegrainedness_split,
      eval_imbalance_dataset,
      omit_from_saving_and_reloading,
      train_episode_config,
      eval_episode_config,
      learn_config,
      learner_config,
      data_config,
  ):
    """Initializes a Trainer.

    Args:
      train_learner: A Learner to be used for meta-training.
      eval_learner: A Learner to be used for meta-validation or meta-testing.
      is_training: Bool, whether or not to train or just evaluate.
      train_dataset_list: A list of names of datasets to train on. This can be
        any subset of the supported datasets.
      eval_dataset_list: A list of names of datasets to evaluate on either for
        validation during train or for final test evaluation, depending on the
        nature of the experiment, as dictated by `is_training'.
      restrict_classes: A dict that maps dataset names to a dict that specifies
        for each of TRAIN_SPLIT, VALID_SPLIT and TEST_SPLIT the number of
        classes to restrict to. This can lead to some classes of a particular
        split of a particular dataset never participating in episode creation.
      restrict_num_per_class: A dict that maps dataset names to a dict that
        specifies for each of TRAIN_SPLIT, VALID_SPLIT and TEST_SPLIT the number
        of examples per class to restrict to. For datasets / splits that are not
        mentioned, no restriction is applied. If restrict_num_per_class is the
        empty dict, no restriction is applied to any split of any dataset.
      checkpoint_dir: A string, the path to the checkpoint directory, or None if
        no checkpointing should occur.
      summary_dir: A string, the path to the checkpoint directory, or None if no
        summaries should be saved.
      records_root_dir: A string, the path to the dataset records directory.
      eval_finegrainedness: Whether to perform binary ImageNet evaluation for
        assessing the performance on fine- vs coarse- grained tasks.
      eval_finegrainedness_split: The subgraph of ImageNet to perform the
        aforementioned analysis on. Notably, if this is TRAIN_SPLIT, we need to
        ensure that an training data is used episodically, even if the given
        model is the baseline model which usually uses batches for training.
      eval_imbalance_dataset: A dataset on which to perform evaluation for
        assessing how class imbalance affects performance in binary episodes. By
        default it is empty and no imbalance analysis is performed.
      omit_from_saving_and_reloading: A list of strings that specifies
        substrings of variable names that should not be reloaded.
      train_episode_config: An instance of EpisodeDescriptionConfig (in
        data/config.py). This is a config for setting the ways and shots of
        training episodes or the parameters for sampling them, if variable.
      eval_episode_config: An instance of EpisodeDescriptionConfig. Analogous to
        train_episode_config but used for eval episodes (validation or testing).
      learn_config: A LearnConfig, the learning configuration.
      learner_config: A LearnerConfig, the learner configuration.
      data_config: A DataConfig, the data configuration.

    Raises:
      UnexpectedSplitError: If split not as expected for Trainer.
    """
    self.train_learner_class = train_learner
    self.eval_learner_class = eval_learner
    self.is_training = is_training
    self.train_dataset_list = train_dataset_list
    self.eval_dataset_list = eval_dataset_list
    self.restrict_classes = restrict_classes
    self.restrict_num_per_class = restrict_num_per_class
    self.checkpoint_dir = checkpoint_dir
    self.summary_dir = summary_dir
    self.records_root_dir = records_root_dir
    self.eval_finegrainedness = eval_finegrainedness
    self.eval_finegrainedness_split = eval_finegrainedness_split
    self.eval_imbalance_dataset = eval_imbalance_dataset
    self.omit_from_saving_and_reloading = omit_from_saving_and_reloading

    self.eval_split = VALID_SPLIT if is_training else TEST_SPLIT
    if eval_finegrainedness:
      # The fine- vs coarse- grained evaluation may potentially be performed on
      # the training graph as it exhibits greater variety in this aspect.
      self.eval_split = eval_finegrainedness_split

    if eval_finegrainedness or eval_imbalance_dataset:
      # We restrict this analysis to the binary classification setting.
      logging.info(
          'Forcing the number of %s classes to be 2, since '
          'the finegrainedness analysis is applied on binary '
          'classification tasks only.', eval_finegrainedness_split)
      if eval_finegrainedness and eval_finegrainedness_split == TRAIN_SPLIT:
        train_episode_config.num_ways = 2
      else:
        eval_episode_config.num_ways = 2

    self.num_train_classes = train_episode_config.num_ways
    self.num_test_classes = eval_episode_config.num_ways
    self.num_support_train = train_episode_config.num_support
    self.num_query_train = train_episode_config.num_query
    self.num_support_eval = eval_episode_config.num_support
    self.num_query_eval = eval_episode_config.num_query

    self.learn_config = learn_config
    self.learner_config = learner_config
    self.train_episode_config = train_episode_config
    self.eval_episode_config = eval_episode_config

    if self.learn_config.transductive_batch_norm:
      logging.warn('Using transductive batch norm!')

    # Only applicable for non-transudctive batch norm. The correct
    # implementation here involves computing the mean and variance based on the
    # support set and then using them to batch normalize the query set. During
    # meta-learning, we allow the gradients to flow through those moments.
    self.backprop_through_moments = True

    self.data_config = data_config
    # Get the image shape.
    self.image_shape = [data_config.image_height] * 2 + [3]

    # Create the benchmark specification.
    self.benchmark_spec = self.get_benchmark_specification()

    # Which splits to support depends on whether we are in the meta-training
    # phase or not. If we are, we need the train split, and the valid one for
    # early-stopping. If not, we only need the test split.
    self.required_splits = [TRAIN_SPLIT] if self.is_training else []
    self.required_splits += [self.eval_split]

    # Get the training, validation and testing specifications.
    # Each is either an EpisodeSpecification or a BatchSpecification.
    self.split_episode_or_batch_specs = {}
    self.next_data = {}
    self.ema_object = None  # Using dummy EMA object for now.
    self.learners = {}
    self.embedding_fn = (
        learner.NAME_TO_EMBEDDING_NETWORK[self.learner_config.embedding_network]
    )
    for split in self.required_splits:
      # Get the next data (episode or batch) for the different splits.
      if split == TRAIN_SPLIT:
        self.split_episode_or_batch_specs[split] = (
            self._create_train_specification())
        # TODO(eringrant): Refactor to expose or avoid dependence of
        # `build_data` on `split_episode_or_batch_specs`.
        self.next_data[split] = self.build_data(split)
        self.learners[split] = self.create_train_learner(
            self.train_learner_class, self.get_next(split))
      elif split in [VALID_SPLIT, TEST_SPLIT]:
        self.split_episode_or_batch_specs[split] = (
            self._create_held_out_specification(split))
        # TODO(eringrant): Refactor to expose or avoid dependence of
        # `build_data` on `split_episode_or_batch_specs`.
        self.next_data[split] = self.build_data(split)
        self.learners[split] = self.create_eval_learner(self.eval_learner_class,
                                                        self.get_next(split))
      else:
        raise UnexpectedSplitError(split)

    # Get the Tensors for the losses / accuracies of the different learners.
    self.losses = dict(
        zip(self.required_splits, [
            self.learners[split].compute_loss()
            for split in self.required_splits
        ]))
    self.accs = dict(
        zip(self.required_splits, [
            self.learners[split].compute_accuracy()
            for split in self.required_splits
        ]))

    # Set self.way, self.shots to Tensors for the way/shots of the next episode.
    self.set_way_shots_classes_logits_targets()

    # Get an optimizer and the operation for meta-training.
    self.train_op = None
    if self.is_training:
      global_step = tf.train.get_or_create_global_step()
      learning_rate = self.learner_config.learning_rate
      if self.learner_config.decay_learning_rate:
        learning_rate = tf.train.exponential_decay(
            self.learner_config.learning_rate,
            global_step,
            decay_steps=self.learner_config.decay_every,
            decay_rate=self.learner_config.decay_rate,
            staircase=True)
      tf.summary.scalar('learning_rate', learning_rate)
      self.optimizer = tf.train.AdamOptimizer(learning_rate)
      self.train_op = self.get_train_op(global_step)

    vars_to_restore = []
    # Omit from reloading any variables that contains as a substring anything in
    # the following list. For example, those that track iterator state, as
    # iterator state is not saved.
    logging.info(
        'Omitting from saving / reloading any variable that '
        'contains any of the following substrings: %s',
        omit_from_saving_and_reloading)
    for var in tf.global_variables():
      if not any([
          substring in var.name for substring in omit_from_saving_and_reloading
      ]):
        vars_to_restore.append(var)
      else:
        logging.info('Omitting variable %s', var.name)
    self.saver = tf.train.Saver(var_list=vars_to_restore, max_to_keep=500)

    if self.checkpoint_dir is not None:
      if not tf.io.gfile.exists(self.checkpoint_dir):
        tf.io.gfile.makedirs(self.checkpoint_dir)

    # Initialize a Session.
    self.initialize_session()
    self.create_summary_writer()

  def set_way_shots_classes_logits_targets(self):
    """Sets the Tensors for the above info of the learner's next episode.

    Raises:
      NotImplementedError: Abstract Method.
    """
    raise NotImplementedError('Abstract Method.')

  def maybe_set_way_shots_classes_logits_targets(self, skip_train=False):
    """Wrapper for set_way_shots_classes_logits_targets.

    Args:
      skip_train: Whether to assign each of the way, shot, (absolute) class id,
        test_logits and test_targets of the training split to None. This is used
        when called with the training split from the batch trainer, since
        training does not happen in episodes there.
    """
    # The batch trainer receives episodes only for the valid and test splits.
    # Therefore for the train split there is no defined way and shots.
    (way, shots, class_props, class_ids, test_logits,
     test_targets) = [], [], [], [], [], []
    for split in self.required_splits:

      if split == TRAIN_SPLIT and skip_train:
        (way_, shots_, class_props_, class_ids_, test_logits_,
         test_targets_) = [None] * 6
      else:
        way_, shots_, class_ids_ = compute_episode_stats(self.next_data[split])
        class_props_ = None
        if self.eval_imbalance_dataset:
          class_props_ = compute_train_class_proportions(
              self.next_data[split], shots_, self.eval_imbalance_dataset_spec)
        test_logits_ = self.learners[split].test_logits
        test_targets_ = self.learners[split].test_targets
      way.append(way_)
      shots.append(shots_)
      class_props.append(class_props_)
      class_ids.append(class_ids_)
      test_logits.append(test_logits_)
      test_targets.append(test_targets_)

    self.way = dict(zip(self.required_splits, way))
    self.shots = dict(zip(self.required_splits, shots))
    self.class_props = dict(zip(self.required_splits, class_props))
    self.class_ids = dict(zip(self.required_splits, class_ids))
    self.test_logits = dict(zip(self.required_splits, test_logits))
    self.test_targets = dict(zip(self.required_splits, test_targets))

  def create_summary_writer(self):
    """Create summaries and writer."""
    # Add summaries for the losses / accuracies of the different learners.
    standard_summaries = []
    for split in self.required_splits:
      loss_summary = tf.summary.scalar('%s_loss' % split, self.losses[split])
      acc_summary = tf.summary.scalar('%s_acc' % split, self.accs[split])
      standard_summaries.append(loss_summary)
      standard_summaries.append(acc_summary)

    # Add summaries for the way / shot / logits / targets of the learners.
    evaluation_summaries = self.add_eval_summaries()

    # All summaries.
    self.standard_summaries = tf.summary.merge(standard_summaries)
    self.evaluation_summaries = tf.summary.merge(evaluation_summaries)

    # Get a writer.
    self.summary_writer = None
    if self.summary_dir is not None:
      self.summary_writer = tf.summary.FileWriter(self.summary_dir)
      if not tf.io.gfile.exists(self.summary_dir):
        tf.io.gfile.makedirs(self.summary_dir)

  def create_train_learner(self, train_learner_class, episode_or_batch):
    """Instantiates a train learner.

    Args:
      train_learner_class: A train learner subclass of the Learner class.
      episode_or_batch: An EpisodeDataset or Batch.

    Raises:
      NotImplementedError: Abstract Method.
    """
    raise NotImplementedError('Abstract Method.')

  def create_eval_learner(self, eval_learner_class, episode):
    """Instantiates an eval learner.

    Args:
      eval_learner_class: An eval learner subclass of the Learner class.
      episode: An EpisodeDataset.

    Raises:
      NotImplementedError: Abstract Method.
    """
    raise NotImplementedError('Abstract Method.')

  def get_benchmark_specification(self, records_root_dir=None):
    """Returns a BenchmarkSpecification.

    Args:
      records_root_dir: Optional. If provided, a list or string that sets the
        directory in which a child directory will be searched for each dataset
        to locate that dataset's records and dataset specification. If it's a
        string, that path will be used for all datasets. If it's a list, its
        length must be the same as the number of datasets, in order to specify a
        different such directory for each. If None, self.records_root_dir will
        be used for all datasets.

    Raises:
      RuntimeError: Incorrect file_pattern detected in a dataset specification.
    """
    (data_spec_list, has_dag_ontology, has_bilevel_ontology,
     splits_to_contribute) = [], [], [], []
    seen_datasets = set()

    if self.is_training:
      benchmark_datasets = self.train_dataset_list + self.eval_dataset_list
    else:
      benchmark_datasets = self.eval_dataset_list

    if isinstance(records_root_dir, list):
      if len(records_root_dir) != len(benchmark_datasets):
        raise ValueError('The given records_root_dir is a list whose length is '
                         'not the same as the number of benchmark datasets. '
                         'Found datasets {} (for the {} phase) but '
                         'len(records_root_dir) is {}. Expected their lengths '
                         'to match or records_path to be a string').format(
                             benchmark_datasets, len(records_root_dir))
      records_roots_for_datasets = records_root_dir
    elif isinstance(records_root_dir, six.text_type):
      records_roots_for_datasets = [records_root_dir] * len(benchmark_datasets)
    elif records_root_dir is None:
      records_roots_for_datasets = [self.records_root_dir
                                   ] * len(benchmark_datasets)

    for dataset_name, dataset_records_root in zip(benchmark_datasets,
                                                  records_roots_for_datasets):

      # Might be seeing a dataset for a second time if it belongs to both the
      # train and eval dataset lists.
      if dataset_name in seen_datasets:
        continue

      dataset_records_path = os.path.join(dataset_records_root, dataset_name)
      data_spec = dataset_spec_lib.load_dataset_spec(dataset_records_path)

      # Only ImageNet has a DAG ontology.
      has_dag = (dataset_name == 'ilsvrc_2012')
      # Only Omniglot has a bi-level ontology.
      is_bilevel = (dataset_name == 'omniglot')

      # The meta-splits that this dataset will contribute data to.
      if not self.is_training:
        # If we're meta-testing, all datasets contribute only to meta-test.
        splits = {self.eval_split}
      else:
        splits = set()
        if dataset_name in self.train_dataset_list:
          splits.add(TRAIN_SPLIT)
        if dataset_name in self.eval_dataset_list:
          splits.add(VALID_SPLIT)

      # By default, all classes of each split will eventually be used for
      # episode creation. But it might be that for some datasets, it is
      # requested to restrict the available number of classes of some splits.
      restricted_classes_per_split = {}
      if dataset_name in self.restrict_classes:
        classes_per_split = self.restrict_classes[dataset_name]
        for split, num_classes in classes_per_split.items():
          # The option to restrict classes is not supported in conjuction with
          # non-uniform (bilevel or hierarhical) class sampling.
          episode_descr_config = (
              self.train_episode_config
              if split == TRAIN_SPLIT else self.eval_episode_config)
          if has_dag and not episode_descr_config.ignore_dag_ontology:
            raise ValueError('Restrictions on the class set of a dataset with '
                             'a DAG ontology are not supported when '
                             'ignore_dag_ontology is False.')
          if is_bilevel and not episode_descr_config.ignore_bilevel_ontology:
            raise ValueError('Restrictions on the class set of a dataset with '
                             'a bilevel ontology are not supported when '
                             'ignore_bilevel_ontology is False.')

          restricted_classes_per_split[get_split_enum(split)] = num_classes
        # Initialize the DatasetSpecificaton to account for this restriction.
        data_spec.initialize(restricted_classes_per_split)

        # Log the applied restrictions.
        tf.logging.info('Restrictions for dataset {}:'.format(dataset_name))
        for split in list(splits):
          num_classes = data_spec.get_classes(get_split_enum(split))
          tf.logging.info('\t split {} is restricted to {} classes'.format(
              split, num_classes))

      # Add this dataset to the benchmark.
      logging.info('Adding dataset %s', data_spec.name)
      data_spec_list.append(data_spec)
      has_dag_ontology.append(has_dag)
      has_bilevel_ontology.append(is_bilevel)
      splits_to_contribute.append(splits)

      # Book-keeping.
      seen_datasets.add(dataset_name)

    if self.eval_imbalance_dataset:
      self.eval_imbalance_dataset_spec = data_spec
      assert len(data_spec_list) == 1, ('Imbalance analysis is only '
                                        'supported on one dataset at a time.')

    benchmark_spec = dataset_spec_lib.BenchmarkSpecification(
        'benchmark', self.image_shape, data_spec_list, has_dag_ontology,
        has_bilevel_ontology, splits_to_contribute)

    # Logging of which datasets will be used for the different meta-splits.
    splits_to_datasets = collections.defaultdict(list)
    for dataset_spec, splits_to_contribute in zip(data_spec_list,
                                                  splits_to_contribute):
      for split in splits_to_contribute:
        splits_to_datasets[split].append(dataset_spec.name)
    for split, datasets in splits_to_datasets.items():
      logging.info('Episodes for split %s will be created from %s', split,
                   datasets)

    return benchmark_spec

  def initialize_session(self):
    """Initializes a tf Session."""
    if ENABLE_TF_OPTIMIZATIONS:
      self.sess = tf.Session()
    else:
      rewriter_config = rewriter_config_pb2.RewriterConfig(
          disable_model_pruning=True,
          constant_folding=rewriter_config_pb2.RewriterConfig.OFF,
          arithmetic_optimization=rewriter_config_pb2.RewriterConfig.OFF,
          remapping=rewriter_config_pb2.RewriterConfig.OFF,
          shape_optimization=rewriter_config_pb2.RewriterConfig.OFF,
          dependency_optimization=rewriter_config_pb2.RewriterConfig.OFF,
          function_optimization=rewriter_config_pb2.RewriterConfig.OFF,
          layout_optimizer=rewriter_config_pb2.RewriterConfig.OFF,
          loop_optimization=rewriter_config_pb2.RewriterConfig.OFF,
          memory_optimization=rewriter_config_pb2.RewriterConfig.NO_MEM_OPT)
      graph_options = tf.GraphOptions(rewrite_options=rewriter_config)
      session_config = tf.ConfigProto(graph_options=graph_options)
      self.sess = tf.Session(config=session_config)

    # Restore or initialize the variables.
    self.sess.run(tf.global_variables_initializer())
    self.sess.run(tf.local_variables_initializer())
    if self.learner_config.checkpoint_for_eval:
      # Requested a specific checkpoint.
      self.saver.restore(self.sess, self.learner_config.checkpoint_for_eval)
      logging.info('Restored checkpoint: %s',
                   self.learner_config.checkpoint_for_eval)
    else:
      # Continue from the latest checkpoint if one exists.
      # This handles fault-tolerance.
      latest_checkpoint = None
      if self.checkpoint_dir is not None:
        latest_checkpoint = tf.train.latest_checkpoint(self.checkpoint_dir)
      if latest_checkpoint:
        self.saver.restore(self.sess, latest_checkpoint)
        logging.info('Restored checkpoint: %s', latest_checkpoint)
      else:
        logging.info('No previous checkpoint.')
        self.sess.run(tf.global_variables_initializer())
        self.sess.run(tf.local_variables_initializer())

    # For episodic models, potentially use pretrained weights at the start of
    # training. If this happens it will overwrite the embedding weights, but
    # taking care to not restore the Adam parameters.
    if self.learner_config.pretrained_checkpoint and not self.sess.run(
        tf.train.get_global_step()):

      # Load the embedding variables from the pre-trained checkpoint. Since the
      # pre-trained checkpoint comes from a BaselineLearner, we need a Saver
      # that only considers embedding Variables from a BaselineLearner. In
      # particular, we exclude 'relationnet*' Variables as they are not present
      # in the checkpoint. We also exclude any variables that are not related
      # to the embedding (e.g. `beta1_power:0') and any variables that are
      # requested to be omitted. Notably, this leads to not reloading ADAM
      # variables. We do not reload these since this episodic finetuning
      # procedure is a different optimization problem than the original training
      # of the baseline whose embedding weights are re-used.
      baselinelearner_embed_vars_to_reload = []
      for var in tf.global_variables():
        is_relationnet_var = var.name.startswith('relationnet')
        requested_to_omit = any([
            substring in var.name
            for substring in self.omit_from_saving_and_reloading
        ])
        is_embedding_var = any(
            keyword in var.name for keyword in EMBEDDING_KEYWORDS)
        is_adam_var = 'adam' in var.name.lower()
        if (not is_relationnet_var and not requested_to_omit and
            is_embedding_var):
          if is_adam_var:
            raise RuntimeError('Variable name unexpectedly indicates it is '
                               'both related to an embedding, and to ADAM: %s' %
                               var.name)
          baselinelearner_embed_vars_to_reload.append(var)
      backbone_saver = tf.train.Saver(
          var_list=baselinelearner_embed_vars_to_reload, max_to_keep=1)
      backbone_saver.restore(self.sess,
                             self.learner_config.pretrained_checkpoint)
      logging.info('Restored only vars %s from checkpoint: %s',
                   [var.name for var in baselinelearner_embed_vars_to_reload],
                   self.learner_config.pretrained_checkpoint)

  def _create_held_out_specification(self, split=TEST_SPLIT):
    """Create an EpisodeSpecification for either validation or testing.

    Note that testing is done episodically whether or not training was episodic.
    This is why the different subclasses should not override this method.

    Args:
      split: one of VALID_SPLIT or TEST_SPLIT

    Returns:
      an EpisodeSpecification.

    Raises:
      ValueError: Invalid split.
    """
    split_enum = get_split_enum(split)
    return learning_spec.EpisodeSpecification(split_enum, self.num_test_classes,
                                              self.num_support_eval,
                                              self.num_query_eval)

  def _create_train_specification(self):
    """Returns an EpisodeSpecification or BatchSpecification for training.

    Raises:
      NotImplementedError: Should be implemented in each subclass.
    """
    raise NotImplementedError('Abstract Method.')

  def build_data(self, split):
    raise NotImplementedError('Abstract method.')

  def _restrict_dataset_list_for_split(self, split, splits_to_contribute,
                                       dataset_list):
    """Returns the restricted dataset_list for the given split.

    Args:
      split: A string, either TRAIN_SPLIT, VALID_SPLIT or TEST_SPLIT.
      splits_to_contribute: A list whose length is the number of datasets in the
        benchmark. Each element is a set of strings corresponding to the splits
        that the respective dataset will contribute to.
      dataset_list: A list which has one element per selected dataset (same
        length as splits_to_contribute), e.g. this can be one of the lists
        dataset_spec_list, has_dag_ontology, has_bilevel_ontology of the
        BenchmarkSpecification.
    """
    updated_list = []
    for dataset_num, dataset_splits in enumerate(splits_to_contribute):
      if split in dataset_splits:
        updated_list.append(dataset_list[dataset_num])
    return updated_list

  def get_num_to_take(self, dataset_name, split):
    """Return the number of examples to restrict to for a dataset/split pair."""
    num_to_take = -1  # By default, no restriction.
    if dataset_name in self.restrict_num_per_class:
      dataset_restrict_num_per_class = self.restrict_num_per_class[dataset_name]
      if split in dataset_restrict_num_per_class:
        num_to_take = dataset_restrict_num_per_class[split]
    return num_to_take

  def build_episode(self, split):
    """Builds an EpisodeDataset containing the next data for "split".

    Args:
      split: A string, either TRAIN_SPLIT, VALID_SPLIT, or TEST_SPLIT.

    Returns:
      An EpisodeDataset.

    Raises:
      UnexpectedSplitError: If split not as expected for this episode build.
    """
    shuffle_buffer_size = self.data_config.shuffle_buffer_size
    read_buffer_size_bytes = self.data_config.read_buffer_size_bytes
    num_prefetch = self.data_config.num_prefetch
    (_, image_shape, dataset_spec_list, has_dag_ontology, has_bilevel_ontology,
     splits_to_contribute) = self.benchmark_spec

    # Choose only the datasets that are chosen to contribute to the given split.
    dataset_spec_list = self._restrict_dataset_list_for_split(
        split, splits_to_contribute, dataset_spec_list)
    has_dag_ontology = self._restrict_dataset_list_for_split(
        split, splits_to_contribute, has_dag_ontology)
    has_bilevel_ontology = self._restrict_dataset_list_for_split(
        split, splits_to_contribute, has_bilevel_ontology)

    episode_spec = self.split_episode_or_batch_specs[split]
    dataset_split = episode_spec[0]
    # TODO(lamblinp): Support non-square shapes if necessary. For now, all
    # images are resized to square, even if it changes the aspect ratio.
    image_size = image_shape[0]
    if image_shape[1] != image_size:
      raise ValueError(
          'Expected a square image shape, not {}'.format(image_shape))

    if split == TRAIN_SPLIT:
      episode_descr_config = self.train_episode_config
    elif split in (VALID_SPLIT, TEST_SPLIT):
      episode_descr_config = self.eval_episode_config
    else:
      raise UnexpectedSplitError(split)

    # Decide how many examples per class to restrict to for each dataset for the
    # given split (by default there is no restriction).
    num_per_class = []  # A list whose length is the number of datasets.
    for dataset_spec in dataset_spec_list:
      num_per_class.append(self.get_num_to_take(dataset_spec.name, split))

    # TODO(lamblinp): pass specs directly to the pipeline builder.
    # TODO(lamblinp): move the special case directly in make_..._pipeline
    if len(dataset_spec_list) == 1:

      use_dag_ontology = has_dag_ontology[0]
      if self.eval_finegrainedness or self.eval_imbalance_dataset:
        use_dag_ontology = False
      data_pipeline = pipeline.make_one_source_episode_pipeline(
          dataset_spec_list[0],
          use_dag_ontology=use_dag_ontology,
          use_bilevel_ontology=has_bilevel_ontology[0],
          split=dataset_split,
          episode_descr_config=episode_descr_config,
          shuffle_buffer_size=shuffle_buffer_size,
          read_buffer_size_bytes=read_buffer_size_bytes,
          num_prefetch=num_prefetch,
          image_size=image_size,
          num_to_take=num_per_class[0])
    else:
      data_pipeline = pipeline.make_multisource_episode_pipeline(
          dataset_spec_list,
          use_dag_ontology_list=has_dag_ontology,
          use_bilevel_ontology_list=has_bilevel_ontology,
          split=dataset_split,
          episode_descr_config=episode_descr_config,
          shuffle_buffer_size=shuffle_buffer_size,
          read_buffer_size_bytes=read_buffer_size_bytes,
          num_prefetch=num_prefetch,
          image_size=image_size,
          num_to_take=num_per_class)

    data_pipeline = apply_dataset_options(data_pipeline)

    iterator = data_pipeline.make_one_shot_iterator()
    (support_images, support_labels, support_class_ids, query_images,
     query_labels, query_class_ids) = iterator.get_next()
    return providers.EpisodeDataset(
        train_images=support_images,
        test_images=query_images,
        train_labels=support_labels,
        test_labels=query_labels,
        train_class_ids=support_class_ids,
        test_class_ids=query_class_ids)

  def build_batch(self, split):
    """Builds a Batch object containing the next data for "split".

    Args:
      split: A string, either TRAIN_SPLIT, VALID_SPLIT, or TEST_SPLIT.

    Returns:
      An EpisodeDataset.
    """
    shuffle_buffer_size = self.data_config.shuffle_buffer_size
    read_buffer_size_bytes = self.data_config.read_buffer_size_bytes
    num_prefetch = self.data_config.num_prefetch
    (_, image_shape, dataset_spec_list, _, _,
     splits_to_contribute) = self.benchmark_spec

    # Choose only the datasets that are chosen to contribute to the given split.
    dataset_spec_list = self._restrict_dataset_list_for_split(
        split, splits_to_contribute, dataset_spec_list)

    # Decide how many examples per class to restrict to for each dataset for the
    # given split (by default there is no restriction).
    num_per_class = []  # A list whose length is the number of datasets.
    for dataset_spec in dataset_spec_list:
      num_per_class.append(self.get_num_to_take(dataset_spec.name, split))

    dataset_split, batch_size = self.split_episode_or_batch_specs[split]
    for dataset_spec in dataset_spec_list:
      if dataset_spec.name in DATASETS_WITH_EXAMPLE_SPLITS:
        raise ValueError(
            'Batch pipeline is used only at meta-train time, and does not '
            'handle datasets with example splits, which should only be used '
            'at meta-test (evaluation) time.')
    # TODO(lamblinp): pass specs directly to the pipeline builder.
    # TODO(lamblinp): move the special case directly in make_..._pipeline
    if len(dataset_spec_list) == 1:
      data_pipeline = pipeline.make_one_source_batch_pipeline(
          dataset_spec_list[0],
          split=dataset_split,
          batch_size=batch_size,
          shuffle_buffer_size=shuffle_buffer_size,
          read_buffer_size_bytes=read_buffer_size_bytes,
          num_prefetch=num_prefetch,
          image_size=image_shape[0],
          num_to_take=num_per_class[0])
    else:
      data_pipeline = pipeline.make_multisource_batch_pipeline(
          dataset_spec_list,
          split=dataset_split,
          batch_size=batch_size,
          shuffle_buffer_size=shuffle_buffer_size,
          read_buffer_size_bytes=read_buffer_size_bytes,
          num_prefetch=num_prefetch,
          image_size=image_shape[0],
          num_to_take=num_per_class)

    data_pipeline = apply_dataset_options(data_pipeline)
    iterator = data_pipeline.make_one_shot_iterator()
    images, class_ids = iterator.get_next()
    return providers.Batch(images=images, labels=class_ids)

  def get_next(self, split):
    """Returns the next batch or episode.

    Args:
      split: A str, one of TRAIN_SPLIT, VALID_SPLIT, or TEST_SPLIT.

    Raises:
      ValueError: Invalid split.
    """
    if split not in [TRAIN_SPLIT, VALID_SPLIT, TEST_SPLIT]:
      raise ValueError('Invalid split. Expected one of "train", "valid", or '
                       '"test".')
    return self.next_data[split]

  def get_train_op(self, global_step):
    """Returns the operation that performs a training update."""
    # UPDATE_OPS picks up batch_norm updates.
    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    with tf.control_dependencies(update_ops):
      train_op = self.optimizer.minimize(
          self.losses[TRAIN_SPLIT], global_step=global_step)
    return train_op

  def get_updated_global_step(self):
    with tf.control_dependencies([self.train_op]):
      global_step = tf.identity(tf.train.get_global_step())
    return global_step

  def train(self):
    """The training loop."""
    global_step = self.sess.run(tf.train.get_global_step())
    logging.info('Starting training from global_step: %d', global_step)
    updated_global_step = self.get_updated_global_step()

    # Dummy variables so that logging works even if called before evaluation.
    self.valid_acc = np.nan
    self.valid_ci = np.nan

    # Compute the initial validation performance before starting the training.
    self.maybe_evaluate(global_step)

    while global_step < self.learn_config.num_updates:
      # Perform the next update.
      (_, train_loss, train_acc, global_step) = self.sess.run([
          self.train_op, self.losses[TRAIN_SPLIT], self.accs[TRAIN_SPLIT],
          updated_global_step
      ])

      # Maybe validate, depending on the global step's value.
      self.maybe_evaluate(global_step)

      # Log training progress.
      if not global_step % self.learn_config.log_every:
        message = (
            'Update %d. Train loss: %f, Train accuracy: %f, '
            'Valid accuracy %f +/- %f.\n' %
            (global_step, train_loss, train_acc, self.valid_acc, self.valid_ci))
        logging.info(message)

        # Update summaries.
        summaries = self.sess.run(self.standard_summaries)
        if self.summary_writer:
          self.summary_writer.add_summary(summaries, global_step)

      should_save = self.checkpoint_dir is not None
      if should_save and global_step % self.learn_config.checkpoint_every == 0:
        save_path = self.saver.save(
            self.sess,
            os.path.join(self.checkpoint_dir, 'model_%d.ckpt' % global_step))
        logging.info('Model checkpoint saved: %s', save_path)

  def maybe_evaluate(self, global_step):
    """Maybe perform evaluation, depending on the value of global_step."""
    if not global_step % self.learn_config.validate_every:
      # Get the validation accuracy and confidence interval.
      (valid_acc, valid_ci, valid_acc_summary,
       valid_ci_summary) = self.evaluate(VALID_SPLIT)
      # Validation summaries are updated every time validation happens which is
      # every validate_every steps instead of log_every steps.
      if self.summary_writer:
        self.summary_writer.add_summary(valid_acc_summary, global_step)
        self.summary_writer.add_summary(valid_ci_summary, global_step)
      self.valid_acc = valid_acc
      self.valid_ci = valid_ci

  def evaluate(self, split):
    """Returns performance metrics across num_eval_trials episodes / batches."""
    num_eval_trials = self.learn_config.num_eval_episodes
    logging.info('Performing evaluation of the %s split using %d episodes...',
                 split, num_eval_trials)
    accuracies = []
    for eval_trial_num in range(num_eval_trials):
      acc, summaries = self.sess.run(
          [self.accs[split], self.evaluation_summaries])
      accuracies.append(acc)
      # Write evaluation summaries.
      if split == self.eval_split and self.summary_writer:
        self.summary_writer.add_summary(summaries, eval_trial_num)
    logging.info('Done.')

    mean_acc = np.mean(accuracies)
    ci_acc = np.std(accuracies) * 1.96 / np.sqrt(len(accuracies))  # confidence

    if split == TEST_SPLIT:
      logging.info('Test accuracy: %f, +/- %f.\n', mean_acc, ci_acc)

    mean_acc_summary = tf.Summary()
    mean_acc_summary.value.add(tag='mean %s acc' % split, simple_value=mean_acc)
    ci_acc_summary = tf.Summary()
    ci_acc_summary.value.add(tag='%s acc CI' % split, simple_value=ci_acc)

    return mean_acc, ci_acc, mean_acc_summary, ci_acc_summary

  def add_eval_summaries_split(self, split):
    """Returns split's summaries of way / shot / classes / logits / targets."""
    split_eval_summaries = []
    way_summary = tf.summary.scalar('%s_way' % split, self.way[split])
    shots_summary = tf.summary.tensor_summary('%s_shots' % split,
                                              self.shots[split])
    classes_summary = tf.summary.tensor_summary('%s_class_ids' % split,
                                                self.class_ids[split])
    logits_summary = tf.summary.tensor_summary('%s_test_logits' % split,
                                               self.test_logits[split])
    targets_summary = tf.summary.tensor_summary('%s_test_targets' % split,
                                                self.test_targets[split])
    if self.eval_imbalance_dataset:
      class_props_summary = tf.summary.tensor_summary('%s_class_props' % split,
                                                      self.class_props[split])
      split_eval_summaries.append(class_props_summary)
    split_eval_summaries.append(way_summary)
    split_eval_summaries.append(shots_summary)
    split_eval_summaries.append(classes_summary)
    split_eval_summaries.append(logits_summary)
    split_eval_summaries.append(targets_summary)
    return split_eval_summaries

class EpisodicTrainer(Trainer):
  """A Trainer that trains a learner through a series of episodes."""

  def build_data(self, split):
    # An EpisodicTrainer will use episodes for all splits.
    return self.build_episode(split)

  def _create_train_specification(self):
    """Returns an EpisodeSpecification or BatchSpecification for training."""
    return learning_spec.EpisodeSpecification(learning_spec.Split.TRAIN,
                                              self.num_train_classes,
                                              self.num_support_train,
                                              self.num_query_train)

  def set_way_shots_classes_logits_targets(self):
    """Sets the Tensors for the above info of the learner's next episode."""
    self.maybe_set_way_shots_classes_logits_targets()

  def add_eval_summaries(self):
    """Returns summaries of way / shot / classes/ logits / targets."""
    evaluation_summaries = []
    for split in self.required_splits:
      evaluation_summaries.extend(self.add_eval_summaries_split(split))
    return evaluation_summaries

  def create_train_learner(self, train_learner_class, episode_or_batch):
    """Instantiates a train learner."""
    return train_learner_class(True, self.learn_config.transductive_batch_norm,
                               self.backprop_through_moments, self.ema_object,
                               self.embedding_fn, episode_or_batch)

  def create_eval_learner(self, eval_learner_class, episode):
    """Instantiates an eval learner."""
    return eval_learner_class(False, self.learn_config.transductive_batch_norm,
                              self.backprop_through_moments, self.ema_object,
                              self.embedding_fn, episode)

class BatchTrainer(Trainer):
  """A Trainer that trains a learner through a series of batches."""

  def build_data(self, split):
    # The TRAIN_SPLIT split is read as batches, VALID_SPLIT and TEST_SPLIT as
    # episodes.
    if split == TRAIN_SPLIT and self.eval_split != TRAIN_SPLIT:
      return self.build_batch(split)
    elif split in (TRAIN_SPLIT, VALID_SPLIT, TEST_SPLIT):
      return self.build_episode(split)
    else:
      raise UnexpectedSplitError(split)

  def _get_num_total_classes(self):
    """Returns total number of classes in the benchmark."""
    # Total class is needed because for ImageNet instead of linearly assigning
    # class id's we use the hierarchy to do this, thus the assumption:
    # {train_ids} < {val_ids} < {test_ids} doesn't hold.
    total_classes = 0
    for dataset_spec in self.benchmark_spec.dataset_spec_list:
      for split in learning_spec.Split:
        total_classes += len(dataset_spec.get_classes(split))
    return total_classes

  def _create_train_specification(self):
    """Returns an EpisodeSpecification or BatchSpecification for training."""
    if self.eval_split == TRAIN_SPLIT:
      return learning_spec.EpisodeSpecification(learning_spec.Split.TRAIN,
                                                self.num_train_classes,
                                                self.num_support_train,
                                                self.num_query_train)
    else:
      return learning_spec.BatchSpecification(learning_spec.Split.TRAIN,
                                              self.learn_config.batch_size)

  def set_way_shots_classes_logits_targets(self):
    """Sets the Tensors for the above info of the learner's next episode."""
    skip_train = True
    if self.eval_split == TRAIN_SPLIT:
      skip_train = False
    self.maybe_set_way_shots_classes_logits_targets(skip_train=skip_train)

  def add_eval_summaries(self):
    """Returns summaries of way / shot / class id's / logits / targets."""
    evaluation_summaries = []
    for split in self.required_splits:
      # In the Batch case, training is non-episodic but evaluation is episodic.
      if split == TRAIN_SPLIT and self.eval_split != TRAIN_SPLIT:
        continue
      evaluation_summaries.extend(self.add_eval_summaries_split(split))
    return evaluation_summaries

  def create_train_learner(self, train_learner_class, episode_or_batch):
    """Instantiates a train learner."""
    num_total_classes = self._get_num_total_classes()
    is_training = False if self.eval_split == TRAIN_SPLIT else True
    return train_learner_class(is_training,
                               self.learn_config.transductive_batch_norm,
                               self.backprop_through_moments, self.ema_object,
                               self.embedding_fn, episode_or_batch,
                               num_total_classes, self.num_test_classes)

  def create_eval_learner(self, eval_learner_class, episode):
    """Instantiates an eval learner."""
    # The eval_learner_class is typically a batch learner, but in the case of
    # the inference-only baselines, it will be an episodic learner. This allows
    # to combine the inference algorithm of Prototypical Networks, etc, as the
    # validation (and test) procedure of a baseline-trained model.
    if eval_learner_class in BATCH_LEARNERS:
      num_total_classes = self._get_num_total_classes()
      return eval_learner_class(False,
                                self.learn_config.transductive_batch_norm,
                                self.backprop_through_moments, self.ema_object,
                                self.embedding_fn, episode, num_total_classes,
                                self.num_test_classes)
    elif eval_learner_class in EPISODIC_LEARNERS:
      return eval_learner_class(False,
                                self.learn_config.transductive_batch_norm,
                                self.backprop_through_moments, self.ema_object,
                                self.embedding_fn, episode)
    else:
      raise ValueError('The specified eval_learner_class should belong to '
                       'BATCH_LEARNERS or EPISODIC_LEARNERS.')

max_stars_repo_name: nik-sm/meta-dataset

max_issues_repo_name: nik-sm/meta-dataset

max_forks_repo_name: nik-sm/meta-dataset


text:
from sqlalchemy import create_engine, Integer, Float, String
import numpy as np
import pandas as pd
import paramiko, sqlalchemy, scp, time, os, re, pathlib, subprocess

def tableIsExist(tablename: str, dbname: str,
                 engine: sqlalchemy.engine.base.Engine) -> bool:
    """
    ===example===
    from sqlalchemy import create_engine
    engine = create_engine("teradata://dbc:dbc@hostip:1025")
    tableIsExist("iris", "db_data", engine)
    >>> True
    """
    table_list = pd.read_sql_query(
        """select * from dbc.TablesV
        where TableName = '{tablename}'
        and DatabaseName = '{dbname}'""".format(
            dbname=dbname, tablename=tablename), engine)
    return True if len(table_list) > 0 else False

def dropIfExists(tablename: str, dbname: str,
                 engine: sqlalchemy.engine.base.Engine) -> None:
    """
    ===example===
    from sqlalchemy import create_engine
    engine = create_engine("teradata://dbc:dbc@hostip:1025")
    dropIfExists("iris", "db_data", engine)
    droped table db_data.iris
    >>> None
    """
    if tableIsExist(tablename, dbname, engine):
        engine.execute(f"drop table {dbname}.{tablename}")
        print(f"droped table {dbname}.{tablename}")

def dtypeParser(series: pd.Series) -> str:
    """
    function to change pandas data type to Teradata data type
    ===example===
    dtypeParser(df["petal_length"])
    >>> "float"
    """
    if series.dtype.name.startswith("int"):
        if series.max() < np.iinfo(np.int32).max and \
                series.min() > np.iinfo(np.int32).min:
            parsed_dtype = "integer"
        elif series.max() < np.iinfo(np.int64).max and \
                series.min() > np.iinfo(np.int64).min:
            parsed_dtype = "bigint"
        else:
            max_length = max(1, max(series.astype(str).str.len())
                             )  # at least 1 length
            parsed_dtype = f"varchar({max_length})"
    elif series.dtype.name.startswith("float64"):
        parsed_dtype = "float"
    else:
        max_length = max(1, max(series.astype(str).str.len())
                         )  # at least 1 length
        parsed_dtype = f"varchar({max_length})"
    return parsed_dtype

def renameColumns(series: pd.Series) -> pd.Series():
    # replace spaces to underscore, delete symbols,
    # add "col_" if name starts with digit,
    # and change "index" to "idx", which comes from reset_index
    res = (series.astype(str).str
           .replace(r"\s", "_", regex=True)
           .replace(r"\W", "", regex=True)
           .replace(r"^(\d)", "col_\\1", regex=True)
           .replace("index", "idx"))
    return res

def createTable(df: pd.DataFrame,
                engine: sqlalchemy.engine.base.Engine,
                tablename: str, dbname: str = None,
                ifExists: str = "error", indexList: list = None,
                isIndexUnique: bool = True, dtype: dict = {}
                ) -> None:
    """
    ===example===
    createTable(df_iris, engine, "iris", "db_data",
                ifExists="replace")
    >>> None
    """
    # make query string
    if dbname is None:
        dbname = engine.url.database
    new_names = renameColumns(df.columns.to_series())
    col_dtype = [dtypeParser(col) if name not in dtype else dtype[name]
                 for name, col in df.iteritems()]

    column_query_part = "\n    ,".join(
        [f"{name} {dtype}" for name, dtype
         in zip(new_names, col_dtype)]
    )
    if indexList:
        idxSeries = pd.Series(indexList).astype(str)
        idx_str = ",".join(renameColumns(idxSeries))
        q_unique = "unique " if isIndexUnique else ""
        index_query_part = q_unique + f"primary index ({idx_str})"
    else:
        index_query_part = "no primary index"
    query = """
    create table {dbname}.{tablename} (
        {column_query_part}
    ) {index_query_part}
    """.format(tablename=tablename, dbname=dbname,
               column_query_part=column_query_part,
               index_query_part=index_query_part)

    # execute query
    if ifExists == "replace":
        dropIfExists(tablename, dbname, engine)
    elif ifExists == "error":
        if tableIsExist(tablename, dbname, engine):
            raise ValueError(f"{tablename}.{dbname} is already exists.")
    elif ifExists == "insert":
        if tableIsExist(tablename, dbname, engine):
            return

    engine.execute(query)
    print("created table \n" + query)

def connectSSH(hostname: str, username: str, password: str = None,
               port: int = 22, keyPath: str = None
               ) -> paramiko.client.SSHClient:
    if keyPath is not None:
        k = paramiko.RSAKey.from_private_key_file(keyPath)
    else:
        k = None
    sshc = paramiko.SSHClient()
    sshc.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    sshc.connect(hostname, port, username, password, pkey=k)
    return sshc

def archiveFile(sourcePath: str, archivePath: str = None, method: str = "gz",
                verbose: bool = True) -> str:
    if archivePath is None:
        archivePath = str(sourcePath) + "." + method

    if method == "7z":
        command = f"7z a {archivePath} {sourcePath}"
    elif method in ["gz", "xz", "bz2"]:
        sourceName = pathlib.Path(sourcePath).name
        sourceFolder = pathlib.Path(sourcePath).parent
        d_format_option = {"gz": "z", "xz": "J", "bz2": "j"}
        command = (f"tar -C {sourceFolder} -{d_format_option[method]}cvf"
                   + f" {archivePath} {sourceName}")
    else:
        raise ValueError(f"method only supports ['7z', 'gz', 'xz', 'bz2']")

    with verbosity_context(f"Archiving {sourcePath} to {archivePath}",
                           verbose):
        if verbose:
            print(subprocess.check_output(command,
                                          universal_newlines=True, shell=True))
        else:
            subprocess.check_output(
                command, universal_newlines=True, shell=True)
    return archivePath

def unarchiveSSH(archivePath: pathlib.PosixPath,
                 sshc: paramiko.client.SSHClient,
                 unarchiveFolder: pathlib.PosixPath = None,
                 method: str = "gz", verbose: bool = True
                 ) -> None:
    if unarchiveFolder is None:
        unarchiveFolder = archivePath.parent

    if method == "7z":
        command = f"7z e {archivePath} -o{unarchiveFolder}"
    elif method in ["gz", "xz", "bz2"]:
        d_format_option = {"gz": "z", "xz": "J", "bz2": "j"}
        command = (f"tar -xv{d_format_option[method]}f"
                   + f"{archivePath} -C {unarchiveFolder}")
    else:
        raise ValueError(f"method only supports ['7z', 'gz', 'xz', 'bz2']")

    with verbosity_context(f"Unarchiving {archivePath}", verbose):
        stdin, stdout, stderr = sshc.exec_command(command)
        if verbose:
            for line in stdout:
                print(line)

def uploadFile(sourcePath: str, targetPath: str,
               sshc: paramiko.client.SSHClient,
               compress_method: str = None,
               verbose: bool = True) -> pathlib.Path:
    def show_progress(filename, size, sent):
        print(f"Uploading {filename} progress: " +
              f"{float(sent)/float(size)*100:.2f}%", end="\r")
    progress = show_progress if verbose else None

    try:
        if compress_method:
            fileName = pathlib.Path(sourcePath).name
            # change targetPath for uploading to
            # targetPath's directory / sourcePath's name + ext.
            targetPath = pathlib.Path(
                str(pathlib.Path(targetPath).parent / fileName) + "."
                + compress_method)
            sourcePath = archiveFile(sourcePath, verbose=verbose,
                                     method=compress_method)
            isArchived = True

        with scp.SCPClient(sshc.get_transport(), progress=progress) as scpc:
            # in case Path is PosixPath, casting them to str
            scpc.put(str(sourcePath), str(targetPath))
            print("\n")  # nextline

        if compress_method:
            unarchiveSSH(targetPath, sshc, method=compress_method,
                         verbose=verbose)
            isUnarchived = True
            # change targetPath to uploaded raw file
            uploadedPath = str(pathlib.Path(targetPath).parent/fileName)
    finally:  # delete archive files
        if 'isArchived' in locals():
            with verbosity_context(f"Deleting archive {sourcePath}",
                                   verbose):
                os.remove(sourcePath)
        if 'isUnarchived' in locals():
            sftp = sshc.open_sftp()
            with verbosity_context(f"Deleting archive {targetPath} via SCP",
                                   verbose):
                sftp.remove(str(targetPath))

    return uploadedPath if compress_method else targetPath

def tdloadViaSSH(engine: sqlalchemy.engine.base.Engine,
                 sshc: paramiko.client.SSHClient,
                 tablename: str, targetPath: str,
                 dbname: str = None,
                 jobname: str = "jobtmp", skipRowNum: int = 0,
                 verbose: bool = True) -> None:

    targetPath = pathlib.Path(targetPath)
    if dbname is None:
        dbname = engine.url.database
    # always use option: QuotedData = 'Optional'
    options = "--DCPQuotedData 'Optional'"
    if skipRowNum > 0:
        options += f" --SourceSkipRows {skipRowNum}"

    tdload_command = (f"tdload -f {targetPath} -t {dbname}.{tablename}"
                      + f" -h {engine.url.host} -u {engine.url.username}"
                      + f" -p {engine.url.password}"
                      + f" --TargetWorkingDatabase {dbname}"
                      + f" {options} {jobname}")

    # drop error log table if exists
    dropIfExists(tablename+"_ET", dbname, engine)
    dropIfExists(tablename+"_UV", dbname, engine)

    # execute command via ssh
    stdin, stdout, stderr = sshc.exec_command(tdload_command)
    for line in stdout:
        if verbose:
            print(line)
        else:
            if re.match(r".*(Total Rows|successfully).*", line):
                print(line)

class verbosity_context:
    def __init__(self, processName: str, verbose: bool = True):
        self.verbose = verbose
        self.processName = processName

    def __enter__(self):
        self.startTime = time.time()
        if self.verbose:
            print(f"start {self.processName}")
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if self.verbose:
            elapsedTime = time.time() - self.startTime
            print(f"    ...finished. : "
                  + f"elapsed time = {elapsedTime:.2f}.sec")

def tdload_df(df: pd.DataFrame, engine: sqlalchemy.engine.base.Engine,
              tablename: str, ssh_ip: str, ssh_username: str,
              dbname: str = None, ifExists: str = "error",
              compress: str = None, dtype: dict = {},
              ssh_password: str = None, ssh_keypath: str = None,
              ssh_folder: str = None,
              dump_folder: str = None,
              saveIndex=False, indexList: list = None,
              isIndexUnique: bool = True, verbose: bool = True) -> None:

    try:
        # 1. save csv
        if dbname is None:
            dbname = engine.url.database
        # use time for avoiding overwrite existing file
        sourceName = f"tmp_{dbname}_{tablename}_{time.time():.0f}.csv"
        if dump_folder is None:
            dump_folder = pathlib.Path().cwd()
        sourcePath = pathlib.Path(dump_folder)/sourceName

        if saveIndex:
            df = df.reset_index()
        with verbosity_context(f"dumping DF to {sourcePath}", verbose):
            csvMade = df.to_csv(sourcePath, index=False)
        # by reading csv again, we can get same string format as csv
        with verbosity_context(f"re-reading {sourcePath}", verbose):
            df = pd.read_csv(sourcePath)

        # 2. create table
        createTable(df, engine, tablename, dbname,
                    ifExists, indexList, isIndexUnique, dtype)

        # 3. copy file with scp
        with verbosity_context(f"connecting ssh", verbose):
            sshc = connectSSH(ssh_ip, ssh_username,
                              ssh_password, keyPath=ssh_keypath)
            isSSHConnected = True
        if ssh_folder is None:
            _, tmp_out, _ = sshc.exec_command('pwd')
            ssh_folder = tmp_out.readline().strip()
        targetPath = pathlib.Path(ssh_folder)/sourceName
        with verbosity_context(f"Uploading File {sourcePath} to {targetPath}",
                               verbose):
            uploadedPath = uploadFile(sourcePath, targetPath, sshc, compress,
                                      verbose)

        # 4. load file with tdload
        with verbosity_context(f"Loading File {uploadedPath} to DB", verbose):
            tdloadViaSSH(engine=engine, sshc=sshc, tablename=tablename,
                         targetPath=uploadedPath, dbname=dbname,
                         skipRowNum=1, verbose=verbose)

    finally:
        # 5. delete tmp files and close connection
        if 'csvMade' in locals():
            with verbosity_context(f"Deleting dumped file {sourcePath}",
                                   verbose):
                os.remove(sourcePath)
        if 'uploadedPath' in locals():
            sftp = sshc.open_sftp()
            with verbosity_context(f"Deleting {uploadedPath} via SCP",
                                   verbose):
                sftp.remove(str(uploadedPath))
        if 'isSSHConnected' in locals():
            with verbosity_context("Closing SSH", verbose):
                sshc.close()

ext: py

reason: import numpy


text:
import cv2
import numpy as np
import csv

from tf_pose.estimator import TfPoseEstimator
from tf_pose.networks import get_graph_path, model_wh

from PyQt5.QtCore import QThread, pyqtSignal

class ModelError(Exception):
    def __init__(self, message):
        super().__init__(message)

class VideoInputError(Exception):
    def __init__(self, message):
        super().__init__(message)

class VideoOutputError(Exception):
    def __init__(self, message):
        super().__init__(message)

class ProcessingInterruptedException(Exception):
    def __init__(self, message):
        super().__init__(message)

class WriteCSVException(Exception):
    def __init__(self, message):
        super(WriteCSVException, self).__init__(message)

class VideoProcesserThread(QThread):

    error_signal = pyqtSignal(Exception)
    finish_signal = pyqtSignal()
    interrupted_signal = pyqtSignal()
    update_signal = pyqtSignal(int)

    def __init__(self, video_path, output_csv, output_video=None, resolution='432x368',
        model="mobilenet_thin", show_bg=True):

        QThread.__init__(self)
        self.video_path = video_path
        self.output_csv = output_csv
        self.output_video = output_video
        self.resolution = resolution
        self.model = model
        self.show_bg = show_bg

        self.is_active = True

    def __del__(self):
        self.wait()

    def __load_model(self, w, h):
        try:
            return TfPoseEstimator(get_graph_path(self.model), target_size=(w, h))
        except Exception:
            raise ModelError(
                "Не удалось загрузить модель: {}. Убедитесь, что модель находится в директории models/graph"
                .format(self.model))

    def __open_csv(self):
        try:
            file_to_save = open(self.output_csv, 'w')
            writer = csv.writer(file_to_save)
            return file_to_save, writer
        except Exception:
            raise WriteCSVException("Не удалось открыть файл csv для записи")

    def __open_video(self):
        try:
            cap = cv2.VideoCapture(self.video_path)
        except Exception:
            raise VideoInputError("Ошибка при чтении видеофайла: \"{}\"".format(self.video_path))
        if not cap.isOpened():
            raise VideoInputError("Ошибка при чтении видеофайла: \"{}\"".format(self.video_path))
        return cap

    def __open_video_writer(self, width, height, fps):
        if self.output_video == '':
            raise VideoOutputError("Не предоставлено путь к файлу вывода")
        try:
            return cv2.VideoWriter(
                self.output_video,
                cv2.VideoWriter_fourcc(*'DIVX'),
                fps,
                (width, height)
            )
        except Exception:
            raise VideoOutputError("Ошибка при создании файла вывода: {}".format(self.output_video))

    def __update_progress(self, frames_total, current_frame):
        percantage = int(float(current_frame) / float(frames_total) * 100)
        self.update_signal.emit(percantage)

    def run(self):
        try:
            w, h = model_wh(self.resolution)
            e = self.__load_model(w, h)
            cap = self.__open_video()

            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            fps = cap.get(cv2.CAP_PROP_FPS)
            frames_total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            current_frame = 0

            csv_file, csv_writer = self.__open_csv()
            video_output = self.__open_video_writer(width, height, fps) if self.output_video else None

            while cap.isOpened():
                if not self.is_active:
                    raise ProcessingInterruptedException("Работа прервана извне")

                ret_val, image = cap.read()
                try:
                    humans = e.inference(image, resize_to_default=(w > 0 and h > 0), upsample_size=4.0)
                except Exception:
                    break

                csv_writer.writerow([humans])

                if not self.show_bg:
                    image = np.zeros(image.shape)
                image = TfPoseEstimator.draw_humans(image, humans, imgcopy=False)

                if video_output:
                    video_output.write(image)

                current_frame += 1
                self.__update_progress(frames_total, current_frame)

            video_output.release()
            csv_file.close()
            cv2.destroyAllWindows()
            self.finish_signal.emit()

        except ProcessingInterruptedException:
            self.interrupted_signal.emit()
        except Exception as e:
            self.error_signal.emit(e)

max_stars_repo_name: AlgieParvin/openpose-gui

max_issues_repo_name: AlgieParvin/openpose-gui

max_forks_repo_name: AlgieParvin/openpose-gui


text:
import numpy
import matplotlib.pyplot as plt
from tqdm import tqdm

r = []

def compare(a, n):
    for i, j in enumerate(a):
        if n < j:
            return i
    return -1

def dokujinokeisannsiki(i, arr):
    return sum(55.81160550 * 0.107**2 / (0.107**2 + (i + arr)**2) **3)

meow = numpy.arange(0, 0.2, 0.001)
for i in tqdm(numpy.arange(-0.2, 0.2, 0.001)):
    if 0 > i:
        r.append(dokujinokeisannsiki(i, meow))
    else:
        ind = compare(meow, i)
        arr = [len(meow[:ind]), len(meow[ind:])]
        nyan = -1*(dokujinokeisannsiki(0, numpy.arange(0, arr[0]*0.001, 0.001)))
        nyan2 = (dokujinokeisannsiki(0, numpy.arange(0, arr[1]*0.001, 0.001)))
        r.append(nyan + nyan2)

plt.plot(numpy.arange(-0.2, 0.2, 0.001), r)
plt.xlabel('position')
plt.ylabel('force')
plt.show()

print(numpy.arange(-0.2, 0.2, 0.001)[numpy.where(numpy.array(r) == numpy.max(numpy.array(r)))])
print(numpy.max(numpy.array(r)))

print(sum(r[:numpy.where(numpy.array(r) == numpy.max(numpy.array(r)))[0][0]]))

max_stars_repo_name: YKunieda/ForPhysics

max_issues_repo_name: YKunieda/ForPhysics

max_forks_repo_name: YKunieda/ForPhysics


text:
from typing import Callable, Iterable, List, Union, Optional

import numpy as np
import pandas as pd
import sklearn.utils.multiclass as multiclass
from category_encoders.utils import convert_input, convert_input_vector
from sklearn.base import BaseEstimator
from sklearn.linear_model import LogisticRegression, Ridge
from sklearn.model_selection import BaseCrossValidator, GridSearchCV

from nyaggle.ensemble.common import EnsembleResult
from nyaggle.validation import cross_validate

def stacking(test_predictions: List[np.ndarray],
             oof_predictions: List[np.ndarray],
             y: pd.Series,
             estimator: Optional[BaseEstimator] = None,
             cv: Optional[Union[int, Iterable, BaseCrossValidator]] = None,
             groups: Optional[pd.Series] = None,
             type_of_target: str = 'auto',
             eval_func: Optional[Callable] = None) -> EnsembleResult:
    """
    Perform stacking on predictions.

    Args:
        test_predictions:
            List of predicted values on test data.
        oof_predictions:
            List of predicted values on out-of-fold training data.
        y:
            Target value
        estimator:
            Estimator used for the 2nd-level model.
            If ``None``, the default estimator (auto-tuned linear model) will be used.
        cv:
            int, cross-validation generator or an iterable which determines the cross-validation splitting strategy.

            - None, to use the default ``KFold(5, random_state=0, shuffle=True)``,
            - integer, to specify the number of folds in a ``(Stratified)KFold``,
            - CV splitter (the instance of ``BaseCrossValidator``),
            - An iterable yielding (train, test) splits as arrays of indices.
        groups:
            Group labels for the samples. Only used in conjunction with a “Group” cv instance (e.g., ``GroupKFold``).
        type_of_target:
            The type of target variable. If ``auto``, type is inferred by ``sklearn.utils.multiclass.type_of_target``.
            Otherwise, ``binary``, ``continuous``, or ``multiclass`` are supported.
        eval_func:
            Evaluation metric used for calculating result score. Used only if ``oof_predictions`` and ``y`` are given.
    Returns:
        Namedtuple with following members

        * test_prediction:
            numpy array, Average prediction on test data.
        * oof_prediction:
            numpy array, Average prediction on Out-of-Fold validation data. ``None`` if ``oof_predictions`` = ``None``.
        * score:
            float, Calculated score on Out-of-Fold data. ``None`` if ``eval_func`` is ``None``.
    """
    assert len(oof_predictions) == len(test_predictions), "Number of oof and test predictions should be same"

    def _stack(predictions):
        if predictions[0].ndim == 1:
            predictions = [p.reshape(len(p), -1) for p in predictions]
        return np.hstack(predictions)

    X_train = convert_input(_stack(oof_predictions))
    y = convert_input_vector(y, X_train.index)
    X_test = convert_input(_stack(test_predictions))

    assert len(X_train) == len(y)

    if type_of_target == 'auto':
        type_of_target = multiclass.type_of_target(y)

    if estimator is None:
        # if estimator is None, tuned linear estimator is used
        if type_of_target == 'continuous':
            estimator = Ridge(normalize=True, random_state=0)
            param_grid = {
                'alpha': [0.001, 0.01, 0.1, 1, 10],
            }
        else:
            estimator = LogisticRegression(random_state=0)
            param_grid = {
                'penalty': ['l1', 'l2'],
                'C': [0.001, 0.01, 0.1, 1, 10],
            }
        grid_search = GridSearchCV(estimator, param_grid, cv=cv)
        grid_search.fit(X_train, y, groups=groups)
        estimator = grid_search.best_estimator_

    result = cross_validate(estimator, X_train, y, X_test, cv=cv, groups=groups, eval_func=eval_func)
    score = result.scores[-1] if result.scores else None

    return EnsembleResult(result.test_prediction, result.oof_prediction, score)

max_stars_repo_name: tenajima/nyaggle

max_issues_repo_name: tenajima/nyaggle

max_forks_repo_name: tenajima/nyaggle


text:
import os
from enum import Enum, auto
from typing import Tuple

import numpy as np
from bokeh.models import Panel, Tabs, Slider, ColumnDataSource, Select, Spinner, Div
from bokeh.plotting import figure
from bokeh.layouts import column, row
from bokeh.server.server import Server

from replay_monitor.db import DBReader

from .monitor import DEFAULT_LOGS_DIR, DEFAULT_DB_FILENAME

def create_state_data_dict_from_state(state: Tuple[np.ndarray]):
    state_data_dicts = []
    for element in state:
        if element.shape[0] == 1:
            element = element[0]
        data = {}
        plot_type = determine_state_element_plot_type(element)
        if plot_type == StateElementPlotType.BAR:
            values = element.flatten()
            data['x'] = list(range(len(values)))
            data['state_element'] = values
            # data['state_x_range_1d'] = range(element.shape[0])
        elif plot_type == StateElementPlotType.MATRIX:
            data['state_element'] = [element]
            # data['state_x_range_2d'] = element.shape[1]
        elif plot_type == StateElementPlotType.COLOR_IMAGE:
            assert element.shape[2] == 3  # Assume color channel is last!
            xdim, ydim = element.shape[0:2]
            img = np.empty((xdim, ydim), dtype=np.uint32)
            view = img.view(dtype=np.uint8).reshape((xdim, ydim, 4))
            view[:, :, :-1] = np.flipud(element)
            view[:, :, -1] = 255
            data['state_element'] = [img]
        state_data_dicts.append(data)
    return state_data_dicts

def create_rewards_data_dict(rewards: np.ndarray):
    return {
        'rewards': rewards,
        'x': list(range(len(data_manager.trajectory_rewards)))
    }

def step_slider_change_handler(attr, old, new):
    global data_manager, ui_manager
    data_manager.change_transition(new)

    ui_manager.step_spinner.value = new
    ui_manager.step_slider.value = new
    ui_manager.update_ui_due_to_transition_change()

def trajectory_changed_handler(attr, old, new):
    global data_manager, ui_manager
    data_manager.change_transition(0, new)

    ui_manager.update_ui_due_to_trajectory_change()

def log_select_change_handler(attr, old, new):
    global data_manager, ui_manager
    data_manager.change_transition(0, 0, new)
    ui_manager.update_ui_due_to_log_change()

class StateElementPlotType(Enum):
    BAR = auto()
    MATRIX = auto()
    COLOR_IMAGE = auto()

def determine_state_element_plot_type(state_elem: np.ndarray) -> StateElementPlotType:
    if len(state_elem.shape) > 1 and state_elem.shape[0] == 1:
        state_elem = state_elem[0]

    if len(state_elem.shape) == 1:
        return StateElementPlotType.BAR
    elif len(state_elem.shape) == 2:
        if 1 in state_elem.shape:
            return StateElementPlotType.BAR
        else:
            return StateElementPlotType.MATRIX
    elif len(state_elem.shape) == 3 and state_elem.shape[2] == 3:
        return StateElementPlotType.COLOR_IMAGE

def create_state_layout(state, data_sources, title):
    tabs_list = []
    for i, element in enumerate(state):
        if len(element.shape) > 1 and element.shape[0] == 1:
            element = element[0]

        plot_type = determine_state_element_plot_type(element)
        if plot_type == StateElementPlotType.BAR:
            fig = figure(plot_width=600, plot_height=600)
            fig.vbar(x='x', width=0.5, bottom=0, top='state_element', source=data_sources[i])
        elif plot_type == StateElementPlotType.MATRIX:
            range_max = max(element.shape[1], element.shape[0])
            fig = figure(plot_width=600, plot_height=600, x_range=(0, range_max), y_range=(0, range_max))
            fig.image(image='state_element', x=0, y=0, dw=element.shape[1], dh=element.shape[0], palette="Spectral11",
                      source=data_sources[i])
        elif plot_type == StateElementPlotType.COLOR_IMAGE:
            range_max = max(element.shape[1], element.shape[0])
            fig = figure(plot_width=600, plot_height=600, x_range=(0, range_max), y_range=(0, range_max))
            fig.image_rgba(image='state_element', x=0, y=0, dw=element.shape[1], dh=element.shape[0],
                           source=data_sources[i])
        else:
            continue

        tabs_list.append(Panel(child=fig, title=f"State Element {i}"))
    state_layout = column(
        Div(text=f"
{title}
"),
        Tabs(tabs=tabs_list)
    )
    return state_layout

class DataManager:
    def __init__(self, db_file_path: str):
        self.db_file_path = db_file_path
        db_reader = DBReader(db_file=db_file_path)
        self.log_ids = db_reader.get_logs_ids()
        self.current_log = self.log_ids[0]
        self.trajectory_index = 0
        self.transition_index = 0
        self.trajectories_lengths = db_reader.get_trajectories_lengths(self.current_log)
        self.n_state_elements = db_reader.get_num_of_state_elements(self.current_log)

        self.trajectory_rewards = db_reader.get_trajectory_rewards(log_id=self.current_log,
                                                                   trajectory_index=self.trajectory_index)

        self.s, self.a, self.r, self.s2 = self.get_transition(self.current_log, self.trajectory_index,
                                                              self.transition_index)

        self.state_elements_sizes = [state_elem.size for state_elem in self.s]

    def get_transition(self, log_id: str, trajectory_index: int, transition_index: int):
        db_reader = DBReader(db_file=self.db_file_path)
        data = db_reader.get_transition_data(log_id, trajectory_index=trajectory_index,
                                             transition_index=transition_index)
        return data

    def change_transition(self, transition_index: int, trajectory_index: int = None, log_id: str = None):
        is_log_changed = False
        if log_id is not None:
            is_log_changed = log_id != self.current_log
            self.current_log = log_id

        if trajectory_index is not None:
            is_different_trajectory = self.trajectory_index != trajectory_index
            self.trajectory_index = trajectory_index
            if is_different_trajectory or is_log_changed:
                db_reader = DBReader(db_file=self.db_file_path)
                self.trajectory_rewards = db_reader.get_trajectory_rewards(log_id=self.current_log,
                                                                           trajectory_index=self.trajectory_index)

        self.transition_index = transition_index

        self.s, self.a, self.r, self.s2 = self.get_transition(self.current_log, self.trajectory_index,
                                                              self.transition_index)

        if log_id is not None and is_log_changed:
            db_reader = DBReader(db_file=self.db_file_path)
            self.trajectories_lengths = db_reader.get_trajectories_lengths(self.current_log)
            self.n_state_elements = db_reader.get_num_of_state_elements(self.current_log)
            self.state_elements_sizes = [state_elem.size for state_elem in self.s]

    def get_current_trajectory_length(self):
        return self.trajectories_lengths[self.trajectory_index]

class UIManager:
    def __init__(self, data_manager: DataManager):
        self.data_manager = data_manager

        self.log_select = Select(title="Choose a log:", value=data_manager.current_log, options=data_manager.log_ids)
        self.log_select.on_change('value', log_select_change_handler)

        self.step_slider = None
        self._generate_step_slider()

        self.trajectory_select = None
        self._generate_trajectory_select()

        self.state_data_sources = None
        self.next_state_data_sources = None
        self.state_layout = None
        self.next_state_layout = None
        self.states_layout = None
        self._create_states_ui()

        self.rewards_data_source = ColumnDataSource(data=create_rewards_data_dict(self.data_manager.trajectory_rewards))
        self.trajectory_rewards_fig = None
        self._create_trajectory_rewards_ui()

        self.textual_action_reward = Div(text=self._generate_textual_action_reward_str())

        self.states_layout = row(self.state_layout, self.next_state_layout)
        self.sliders_layout = column(row(self.step_slider, self.step_spinner), self.trajectory_select, sizing_mode='stretch_width')
        self.layout = column(self.log_select,
                             self.trajectory_rewards_fig,
                             self.textual_action_reward,
                             self.states_layout,
                             self.sliders_layout)

    def update_ui_due_to_log_change(self):
        self._create_states_ui()
        self.states_layout.children = [self.state_layout, self.next_state_layout]

        self.rewards_data_source.stream(new_data=create_rewards_data_dict(self.data_manager.trajectory_rewards),
                                        rollover=len(self.data_manager.trajectory_rewards))

        self.trajectory_select.high = len(self.data_manager.trajectories_lengths)-1
        self.trajectory_select.value = self.data_manager.trajectory_index

        self.step_slider.value = self.data_manager.transition_index
        self.step_slider.end = self.data_manager.get_current_trajectory_length()-1
        self.step_spinner.high = self.step_slider.end
        self.step_spinner.value = self.step_slider.value

    def update_ui_due_to_transition_change(self):
        state_data_dicts = create_state_data_dict_from_state(self.data_manager.s)
        next_state_data_dicts = create_state_data_dict_from_state(self.data_manager.s2)
        for i, element in enumerate(self.data_manager.s):
            self.state_data_sources[i].stream(new_data=state_data_dicts[i],
                                                    rollover=self.data_manager.state_elements_sizes[i])
            self.next_state_data_sources[i].stream(new_data=next_state_data_dicts[i],
                                                         rollover=self.data_manager.state_elements_sizes[i])

        self.textual_action_reward.text = self._generate_textual_action_reward_str()

    def update_ui_due_to_trajectory_change(self):
        self.update_ui_due_to_transition_change()

        self.rewards_data_source.stream(new_data=create_rewards_data_dict(self.data_manager.trajectory_rewards),
                                        rollover=len(self.data_manager.trajectory_rewards))

        self.step_slider.value = self.data_manager.transition_index
        self.step_slider.end = self.data_manager.get_current_trajectory_length() - 1
        self.step_spinner.high = self.step_slider.end
        self.step_spinner.value = self.step_slider.value

    def _generate_step_slider(self):
        self.step_slider = Slider(start=0, end=self.data_manager.get_current_trajectory_length()-1,
                                  value=self.data_manager.transition_index, step=1, title="Time Step",
                                  sizing_mode='stretch_width')
        self.step_spinner = Spinner(low=0, high=self.step_slider.end, step=1, value=self.step_slider.value, width=100)
        self.step_slider.on_change('value', step_slider_change_handler)
        self.step_spinner.on_change('value', step_slider_change_handler)

    def _generate_trajectory_select(self):
        self.trajectory_select = Spinner(
            title="Choose a trajectory:",
            low=0,
            high=len(self.data_manager.trajectories_lengths)-1,
            step=1,
            value=self.data_manager.trajectory_index,
        )
        self.trajectory_select.on_change('value', trajectory_changed_handler)

    def _create_states_ui(self):
        self.state_data_sources = [ColumnDataSource(data=data_dict)
                                   for data_dict in create_state_data_dict_from_state(self.data_manager.s)]
        self.next_state_data_sources = [ColumnDataSource(data=data_dict)
                                        for data_dict in create_state_data_dict_from_state(self.data_manager.s2)]
        self.state_layout = create_state_layout(self.data_manager.s, self.state_data_sources, 'State')
        self.next_state_layout = create_state_layout(self.data_manager.s2, self.next_state_data_sources, 'Next State')

    def _create_trajectory_rewards_ui(self):
        fig = figure(plot_width=600, plot_height=300, sizing_mode='stretch_width')
        fig.circle(x='x', y='rewards', source=self.rewards_data_source, name='trajectory_rewards_plot')
        self.trajectory_rewards_fig = fig
        return fig

    def _generate_textual_action_reward_str(self):
        text = f'
Reward: {self.data_manager.r[0]}
Action: {self.data_manager.a[0]}
'
        return text

db_file_path = os.path.join(DEFAULT_LOGS_DIR, DEFAULT_DB_FILENAME)
data_manager = None
ui_manager = None

def start_app(doc):
    global data_manager, ui_manager, db_file_path
    data_manager = DataManager(db_file_path)
    ui_manager = UIManager(data_manager=data_manager)
    doc.add_root(ui_manager.layout)

def _start_server(db_path: str = None):
    global db_file_path

    if db_path is not None:
        db_file_path = db_path

    server = Server({'/': start_app}, num_procs=1)
    server.start()

    print('Opening Bokeh application on http://localhost:5006/')

    server.io_loop.add_callback(server.show, "/")
    server.io_loop.start()

def start_server():
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('--db_path', required=False)
    args = parser.parse_args()

    _start_server(args.db_path)

if __name__ == '__main__':
    start_server()


text:
# bandit.generate
# Create a graph from a specified topology.
#
# Author:  Benjamin Bengfort <benjamin@bengfort.com>
# Created: Mon Jun 25 11:08:25 2018 -0400
#
# ID: generate.py [] benjamin@bengfort.com $

"""
Create a graph from a specified topology.
"""

##########################################################################
## Imports
##########################################################################

import csv
import argparse
import networkx as nx

from .serialize import write_graph, CHOICES

def graph_from_latency(path):
    G = nx.DiGraph(name="Network Topology")
    with open(path, 'r') as f:
        reader = csv.DictReader(f)
        for row in reader:
            data = {key: float(row[key]) for key in ('pull', 'push', 'sync')}

            G.add_edge(row['src'], row['dst'], weight=data['sync'], **data)

    return G

##########################################################################
## Main
##########################################################################

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="generate network graph from results file"
    )

    parser.add_argument('-f', '--format', choices=CHOICES, default='json', help='graph serialization format')
    parser.add_argument('-o', '--outpath', default='topology.json', help='path to write graph to')
    parser.add_argument('latency', nargs=1, help='latency.csv file from results')

    args = parser.parse_args()
    G = graph_from_latency(args.latency[0])
    print(nx.info(G))

    write_graph(G, args.outpath, args.format)
    print("{} written to {}".format(args.format, args.outpath))

max_stars_repo_name: bbengfort/anti-entropy-bandits

max_issues_repo_name: bbengfort/anti-entropy-bandits

max_forks_repo_name: bbengfort/anti-entropy-bandits

reason: import networkx


text:
import numpy as np
from sklearn.externals import joblib

# These are the feature labels from our data set
feature_labels = np.array(['year_built', 'stories', 'num_bedrooms', 'full_bathrooms', 'half_bathrooms', 'livable_sqft', 'total_sqft', 'garage_sqft', 'carport_sqft', 'has_fireplace', 'has_pool', 'has_central_heating', 'has_central_cooling', 'garage_type_attached', 'garage_type_detached', 'garage_type_none', 'city_Amystad', 'city_Brownport', 'city_Chadstad', 'city_Clarkberg', 'city_Coletown', 'city_Davidfort', 'city_Davidtown', 'city_East Amychester', 'city_East Janiceville', 'city_East Justin', 'city_East Lucas', 'city_Fosterberg', 'city_Hallfort', 'city_Jeffreyhaven', 'city_Jenniferberg', 'city_Joshuafurt', 'city_Julieberg', 'city_Justinport', 'city_Lake Carolyn', 'city_Lake Christinaport', 'city_Lake Dariusborough', 'city_Lake Jack', 'city_Lake Jennifer', 'city_Leahview', 'city_Lewishaven', 'city_Martinezfort', 'city_Morrisport', 'city_New Michele', 'city_New Robinton', 'city_North Erinville', 'city_Port Adamtown', 'city_Port Andrealand', 'city_Port Daniel', 'city_Port Jonathanborough', 'city_Richardport', 'city_Rickytown', 'city_Scottberg', 'city_South Anthony', 'city_South Stevenfurt', 'city_Toddshire', 'city_Wendybury', 'city_West Ann', 'city_West Brittanyview', 'city_West Gerald', 'city_West Gregoryview', 'city_West Lydia', 'city_West Terrence'])

# Load the trained model created with train_model.py
model = joblib.load('trained_house_classifier_model.pkl')

# Create a numpy array based on the model's feature importances
importance = model.feature_importances_

# Sort the feature labels based on the feature importance rankings from the model
feauture_indexes_by_importance = importance.argsort()

# Print each feature label, from most important to least important (reverse order)
for index in feauture_indexes_by_importance:
    print("{} - {:.2f}%".format(feature_labels[index], (importance[index] * 100.0)))

max_stars_repo_path: Chapter 6/feature_selection final.py

max_issues_repo_path: Chapter 6/feature_selection final.py

max_forks_repo_path: Chapter 6/feature_selection final.py