train_app.py

# vim: expandtab:ts=4:sw=4
import argparse
import numpy as np
import tensorflow as tf
import tensorflow.contrib.slim as slim

from datasets import util
import queued_trainer
import metrics
import losses


def create_default_argument_parser(dataset_name):
    """Create an argument parser with default arguments.

    Parameters
    ----------
    dataset_name : str
        Name of the dataset. This value is used to set default directories.

    Returns
    -------
    argparse.ArgumentParser
        Returns an argument parser with default arguments.

    """
    parser = argparse.ArgumentParser(
        description="Metric trainer (%s)" % dataset_name)
    parser.add_argument(
        "--batch_size", help="Training batch size", default=128, type=int)
    parser.add_argument(
        "--learning_rate", help="Learning rate", default=1e-3, type=float)
    parser.add_argument(
        "--eval_log_dir",
        help="Evaluation log directory (only used in mode 'evaluation').",
        default="/tmp/%s_evaldir" % dataset_name)
    parser.add_argument(
        "--number_of_steps", help="Number of train/eval steps. If None given, "
        "runs indefenitely", default=None, type=int)
    parser.add_argument(
        "--log_dir", help="Log and checkpoints directory.",
        default="/tmp/%s_logdir" % dataset_name)
    parser.add_argument(
        "--loss_mode", help="One of 'cosine-softmax', 'magnet', 'triplet'",
        type=str, default="cosine-softmax")
    parser.add_argument(
        "--mode", help="One of 'train', 'eval', 'finalize', 'freeze'.",
        type=str, default="train")
    parser.add_argument(
        "--restore_path", help="If not None, resume training of a given "
        "checkpoint (mode 'train').", default=None)
    parser.add_argument(
        "--run_id", help="An optional run-id. If None given, a new one is "
        "created", type=str, default=None)
    return parser


def to_train_kwargs(args):
    """Parse command-line training arguments.

    Parameters
    ----------
    args : argparse.Namespace
        Namespace of an argument parser that was created with
        create_default_argument_parser.

    Returns
    -------
    Dict[str, T]
        Returns a dictionary of named arguments to be passed on to
        train_loop.

    """
    kwargs_dict = {
        "batch_size": args.batch_size,
        "learning_rate": args.learning_rate,
        "log_dir": args.log_dir,
        "loss_mode": args.loss_mode,
        "number_of_steps": args.number_of_steps,
        "restore_path": args.restore_path,
        "run_id": args.run_id,
    }
    return kwargs_dict


def to_eval_kwargs(args):
    """Parse command-line evaluation arguments.

    Parameters
    ----------
    args : argparse.Namespace
        Namespace of an argument parser that was created with
        create_default_argument_parser.

    Returns
    -------
    Dict[str, T]
        Returns a dictionary of named arguments to be passed on to
        eval_loop.

    """
    kwargs_dict = {
        "eval_log_dir": args.eval_log_dir,
        "log_dir": args.log_dir,
        "loss_mode": args.loss_mode,
        "run_id": args.run_id,
    }
    return kwargs_dict


def train_loop(preprocess_fn, network_factory, train_x, train_y,
               num_images_per_id, batch_size, log_dir, image_shape=None,
               restore_path=None, exclude_from_restore=None, run_id=None,
               number_of_steps=None, loss_mode="cosine-softmax",
               learning_rate=1e-3, trainable_scopes=None,
               save_summaries_secs=60, save_interval_secs=300):
    """Start training.

    Parameters
    ----------
    preprocess_fn : Callable[tf.Tensor] -> tf.Tensor
        A callable that applies preprocessing to a given input image tensor of
        dtype tf.uint8 and returns a floating point representation (tf.float32).
    network_factory : Callable[tf.Tensor] -> (tf.Tensor, tf.Tensor)
        A callable that takes as argument a preprocessed input image of dtype
        tf.float32 and returns the feature representation as well as a logits
        tensors. The logits may be set to None if not required by the loss.
    train_x : List[str] | np.ndarray
        A list of image filenames or a tensor of images.
    train_y : List[int] | np.ndarray
        A list or one-dimensional array of labels for the images in `train_x`.
    num_images_per_id : int
        Sample `num_images_per_id` images for each label at each training
        iteration. The number of identities sampled at each iteration is
        computed as `batch_size / num_images_per_id`. The `batch_size` must be
        divisible by this number.
    batch_size : int
        The number of images at each training iteration.
    log_dir : str
        Used to construct the log and checkpoint directory. They are stored in
        `log_dir/run_id`.
    image_shape : Tuple[int, int, int] | NoneType
        Image shape (height, width, channels) or None. If None, `train_x` must
        be an array of images such that the shape can be queries from this
        variable.
    restore_path : Optional[str]
        If not None, resumes training from the given checkpoint file.
    exclude_from_restore : Optional[List[str]]
        An optional list of variable scopes to be used in conjunction with
        `restore_path`. If not None, variables in the given scopes are not
        restored from the checkpoint file.
    run_id : Optional[str]
        A string that identifies the training run; used to construct the
        log and checkpoint directory `log_dir/run_id`. If None, a random
        string is created.
    number_of_steps : Optional[int]
        The total number of training iterations. If None, training runs
        indefenitely.
    loss_mode : Optional[str]
        A string that identifies the loss function used for training; must be
        one of 'cosine-softmax', 'magnet', 'triplet'. This value defaults to
        'cosine-softmax'.
    learning_rate : Optional[float]
        Adam learning rate; defaults to 1e-3.
    trainable_scopes : Optional[List[str]]
        Optional list of variable scopes. If not None, only variables within the
        given scopes are trained. Otherwise all variables are trained.
    save_summaries_secs : Optional[int]
        Save training summaries every `save_summaries_secs` seconds to the
        log directory.
    save_interval_secs : Optional[int]
        Save checkpoints every `save_interval_secs` seconds to the log
        directory.

    """
    if image_shape is None:
        # If image_shape is not set, train_x must be an image array. Here we
        # query the image shape from the array of images.
        assert type(train_x) == np.ndarray
        image_shape = train_x.shape[1:]
    elif type(train_x) == np.ndarray:
        assert train_x.shape[1:] == image_shape
    read_from_file = type(train_x) != np.ndarray

    trainer, train_op = create_trainer(
        preprocess_fn, network_factory, read_from_file, image_shape, batch_size,
        loss_mode, learning_rate=learning_rate,
        trainable_scopes=trainable_scopes)

    feed_generator = queued_trainer.random_sample_identities_forever(
        batch_size, num_images_per_id, train_x, train_y)

    variables_to_restore = slim.get_variables_to_restore(
        exclude=exclude_from_restore)
    trainer.run(
        feed_generator, train_op, log_dir, restore_path=restore_path,
        variables_to_restore=variables_to_restore,
        run_id=run_id, save_summaries_secs=save_summaries_secs,
        save_interval_secs=save_interval_secs, number_of_steps=number_of_steps)


def create_trainer(preprocess_fn, network_factory, read_from_file, image_shape,
                   batch_size, loss_mode, learning_rate=1e-3,
                   trainable_scopes=None):
    """Create trainer.

    Parameters
    ----------
    preprocess_fn : Callable[tf.Tensor] -> tf.Tensor
        A callable that applies preprocessing to a given input image tensor of
        dtype tf.uint8 and returns a floating point representation (tf.float32).
    network_factory : Callable[tf.Tensor] -> (tf.Tensor, tf.Tensor)
        A callable that takes as argument a preprocessed input image of dtype
        tf.float32 and returns the feature representation as well as a logits
        tensors. The logits may be set to None if not required by the loss.
    read_from_file:
        Set to True if images are read from file. If False, the trainer expects
        input images as numpy arrays (i.e., data loading must be handled outside
        of the trainer).
    image_shape: Tuple[int, int, int]
        Image shape (height, width, channels).
    batch_size:
        Number of images per batch.
    loss_mode : str
        One of 'cosine-softmax', 'magnet', 'triplet'. If 'cosine-softmax', the
        logits tensor returned by the `network_factory` must not be None.
    learning_rate: float
        Adam learning rate; defauls to 1e-3.
    trainable_scopes: Optional[List[str]]
        Optional list of variable scopes. If not None, only variables within the
        given scopes are trained. Otherwise all variables are trained.

    Returns
    -------
    QueuedTrainer
        Returns a trainer object to be used for training and evaluating the
        given TensorFlow model.

    """
    num_channels = image_shape[-1] if len(image_shape) == 3 else 1

    with tf.device("/cpu:0"):
        label_var = tf.placeholder(tf.int64, (None,))

        if read_from_file:
            # NOTE(nwojke): tf.image.decode_jpg handles various image types.
            filename_var = tf.placeholder(tf.string, (None, ))
            image_var = tf.map_fn(
                lambda x: tf.image.decode_jpeg(
                    tf.read_file(x), channels=num_channels),
                filename_var, back_prop=False, dtype=tf.uint8)
            image_var = tf.image.resize_images(image_var, image_shape[:2])
            input_vars = [filename_var, label_var]
        else:
            image_var = tf.placeholder(tf.uint8, (None,) + image_shape)
            input_vars = [image_var, label_var]

        enqueue_vars = [
            tf.map_fn(
                lambda x: preprocess_fn(x, is_training=True),
                image_var, back_prop=False, dtype=tf.float32),
            label_var]

    trainer = queued_trainer.QueuedTrainer(enqueue_vars, input_vars)
    image_var, label_var = trainer.get_input_vars(batch_size)
    tf.summary.image("images", image_var)

    feature_var, logit_var = network_factory(image_var)
    _create_loss(feature_var, logit_var, label_var, mode=loss_mode)

    if trainable_scopes is None:
        variables_to_train = tf.trainable_variables()
    else:
        variables_to_train = []
        for scope in trainable_scopes:
            variables = tf.get_collection(
                tf.GraphKeys.TRAINABLE_VARIABLES, scope)
            variables_to_train.extend(variables)

    global_step = tf.train.get_or_create_global_step()

    loss_var = tf.losses.get_total_loss()
    train_op = slim.learning.create_train_op(
        loss_var, tf.train.AdamOptimizer(learning_rate=learning_rate),
        global_step, summarize_gradients=False,
        variables_to_train=variables_to_train)
    tf.summary.scalar("total_loss", loss_var)
    tf.summary.scalar("learning_rate", learning_rate)

    regularization_var = tf.reduce_sum(tf.losses.get_regularization_loss())
    tf.summary.scalar("weight_loss", regularization_var)
    return trainer, train_op


def eval_loop(preprocess_fn, network_factory, data_x, data_y, camera_indices,
              log_dir, eval_log_dir, image_shape=None, run_id=None,
              loss_mode="cosine-softmax", num_galleries=10, random_seed=4321):
    """Evaluate a running training session using CMC metric averaged over
    `num_galleries` galleries where each gallery contains for every identity a
    randomly selected image-pair.

    A call to this function will block indefinitely, monitoring the
    `log_dir/run_id` for saved checkpoints. Then, creates summaries in
    `eval_log_dir/run_id` that can be monitored with tensorboard.

    Parameters
    ----------
    preprocess_fn : Callable[tf.Tensor] -> tf.Tensor
        A callable that applies preprocessing to a given input image tensor of
        dtype tf.uint8 and returns a floating point representation (tf.float32).
    network_factory : Callable[tf.Tensor] -> (tf.Tensor, tf.Tensor)
        A callable that takes as argument a preprocessed input image of dtype
        tf.float32 and returns the feature representation as well as a logits
        tensors. The logits may be set to None if not required by the loss.
    data_x : List[str] | np.ndarray
        A list of image filenames or a tensor of images.
    data_y : List[int] | np.ndarray
        A list or one-dimensional array of labels for the images in `data_x`.
    camera_indices: Optional[List[int] | np.ndarray]
        A list or one-dimensional array of camera indices for the images in
        `data_x`. If not None, CMC galleries are created such that image pairs
        are collected from different cameras.
    log_dir: str
        Should be equivalent to the `log_dir` passed into `train_loop` of the
        training run to monitor.
    eval_log_dir:
        Used to construct the tensorboard log directory where metrics are
        summarized.
    image_shape : Tuple[int, int, int] | NoneType
        Image shape (height, width, channels) or None. If None, `train_x` must
        be an array of images such that the shape can be queries from this
        variable.
    run_id : str
        A string that identifies the training run; must be set to the same
        `run_id` passed into `train_loop`.
    loss_mode : Optional[str]
        A string that identifies the loss function used for training; must be
        one of 'cosine-softmax', 'magnet', 'triplet'. This value defaults to
        'cosine-softmax'.
    num_galleries: int
        The number of galleries to be constructed for evaluation of CMC
        metrics.
    random_seed: Optional[int]
        If not None, the NumPy random seed is fixed to this number; can be used
        to produce the same galleries over multiple runs.

    """
    if image_shape is None:
        # If image_shape is not set, train_x must be an image array. Here we
        # query the image shape from the array of images.
        assert type(data_x) == np.ndarray
        image_shape = data_x.shape[1:]
    elif type(data_x) == np.ndarray:
        assert data_x.shape[1:] == image_shape
    read_from_file = type(data_x) != np.ndarray

    # Create num_galleries random CMC galleries to average CMC top-k over.
    probes, galleries = [], []
    for i in range(num_galleries):
        probe_indices, gallery_indices = util.create_cmc_probe_and_gallery(
            data_y, camera_indices, seed=random_seed + i)
        probes.append(probe_indices)
        galleries.append(gallery_indices)
    probes, galleries = np.asarray(probes), np.asarray(galleries)

    # Set up the data feed.
    with tf.device("/cpu:0"):
        # Feed probe and gallery indices to the trainer.
        num_probes, num_gallery_images = probes.shape[1], galleries.shape[1]
        probe_idx_var = tf.placeholder(tf.int64, (None, num_probes))
        gallery_idx_var = tf.placeholder(tf.int64, (None, num_gallery_images))
        trainer = queued_trainer.QueuedTrainer(
            [probe_idx_var, gallery_idx_var])

        # Retrieve indices from trainer and gather data from constant memory.
        data_x_var = tf.constant(data_x)
        data_y_var = tf.constant(data_y)

        probe_idx_var, gallery_idx_var = trainer.get_input_vars(batch_size=1)
        probe_idx_var = tf.squeeze(probe_idx_var)
        gallery_idx_var = tf.squeeze(gallery_idx_var)

        # Apply preprocessing.
        probe_x_var = tf.gather(data_x_var, probe_idx_var)
        if read_from_file:
            # NOTE(nwojke): tf.image.decode_jpg handles various image types.
            num_channels = image_shape[-1] if len(image_shape) == 3 else 1
            probe_x_var = tf.map_fn(
                lambda x: tf.image.decode_jpeg(
                    tf.read_file(x), channels=num_channels),
                probe_x_var, dtype=tf.uint8)
            probe_x_var = tf.image.resize_images(probe_x_var, image_shape[:2])
        probe_x_var = tf.map_fn(
            lambda x: preprocess_fn(x, is_training=False),
            probe_x_var, back_prop=False, dtype=tf.float32)
        probe_y_var = tf.gather(data_y_var, probe_idx_var)

        gallery_x_var = tf.gather(data_x_var, gallery_idx_var)
        if read_from_file:
            # NOTE(nwojke): tf.image.decode_jpg handles various image types.
            num_channels = image_shape[-1] if len(image_shape) == 3 else 1
            gallery_x_var = tf.map_fn(
                lambda x: tf.image.decode_jpeg(
                    tf.read_file(x), channels=num_channels),
                gallery_x_var, dtype=tf.uint8)
            gallery_x_var = tf.image.resize_images(
                gallery_x_var, image_shape[:2])
        gallery_x_var = tf.map_fn(
            lambda x: preprocess_fn(x, is_training=False),
            gallery_x_var, back_prop=False, dtype=tf.float32)
        gallery_y_var = tf.gather(data_y_var, gallery_idx_var)

    # Construct the network and compute features.
    probe_and_gallery_x_var = tf.concat(
        axis=0, values=[probe_x_var, gallery_x_var])
    probe_and_gallery_x_var, _ = network_factory(probe_and_gallery_x_var)

    num_probe = tf.shape(probe_x_var)[0]
    probe_x_var = tf.slice(
        probe_and_gallery_x_var, [0, 0], [num_probe, -1])
    gallery_x_var = tf.slice(
        probe_and_gallery_x_var, [num_probe, 0], [-1, -1])

    # Set up the metrics.
    distance_measure = (
        metrics.cosine_distance if loss_mode == "cosine-softmax"
        else metrics.pdist)

    def cmc_metric_at_k(k):
        return metrics.streaming_mean_cmc_at_k(
            probe_x_var, probe_y_var, gallery_x_var, gallery_y_var,
            k=k, measure=distance_measure)

    names_to_values, names_to_updates = slim.metrics.aggregate_metric_map({
        "Precision@%d" % k: cmc_metric_at_k(k) for k in [1, 5, 10, 20]})
    for metric_name, metric_value in names_to_values.items():
        tf.summary.scalar(metric_name, metric_value)

    # Start evaluation loop.
    trainer.evaluate(
        (probes, galleries), log_dir, eval_log_dir, run_id=run_id,
        eval_op=list(names_to_updates.values()), eval_interval_secs=60)


def finalize(preprocess_fn, network_factory, checkpoint_path, image_shape,
             output_filename):
    """Finalize model, i.e., strip off training variables and only save model
    variables to checkpoint file.

    Parameters
    ----------
    preprocess_fn : Callable[tf.Tensor] -> tf.Tensor
        A callable that applies preprocessing to a given input image tensor of
        dtype tf.uint8 and returns a floating point representation (tf.float32).
    network_factory : Callable[tf.Tensor] -> (tf.Tensor, tf.Tensor)
        A callable that takes as argument a preprocessed input image of dtype
        tf.float32 and returns the feature representation as well as a logits
        tensors. The logits may be set to None if not required by the loss.
    checkpoint_path : str
        The checkpoint file to load.
    image_shape : Tuple[int, int, int]
        Image shape (height, width, channels).
    output_filename : str
        The checkpoint file to write.

    """
    with tf.Session(graph=tf.Graph()) as session:
        input_var = tf.placeholder(tf.uint8, (None, ) + image_shape)
        image_var = tf.map_fn(
            lambda x: preprocess_fn(x, is_training=False),
            input_var, back_prop=False, dtype=tf.float32)
        network_factory(image_var)

        loader = tf.train.Saver(slim.get_variables_to_restore())
        loader.restore(session, checkpoint_path)

        saver = tf.train.Saver(slim.get_model_variables())
        saver.save(session, output_filename, global_step=None)


def freeze(preprocess_fn, network_factory, checkpoint_path, image_shape,
           output_filename, input_name="images", feature_name="features"):
    """Write frozen inference graph that takes as input a list of images and
    returns their feature representation.

    Parameters
    ----------
    preprocess_fn : Callable[tf.Tensor] -> tf.Tensor
        A callable that applies preprocessing to a given input image tensor of
        dtype tf.uint8 and returns a floating point representation (tf.float32).
    network_factory : Callable[tf.Tensor] -> (tf.Tensor, tf.Tensor)
        A callable that takes as argument a preprocessed input image of dtype
        tf.float32 and returns the feature representation as well as a logits
        tensors. The logits may be set to None if not required by the loss.
    checkpoint_path : str
        The checkpoint file to load.
    image_shape : Tuple[int, int, int]
        Image shape (height, width, channels).
    output_filename : str
        Path to the file to write to.
    input_name : Optional[str]
        The input (image) placeholder will be given this name; defaults
        to `images`.
    feature_name : Optional[str]
        The output (feature) tensor will be given this name; defaults to
        `features`.

    """
    with tf.Session(graph=tf.Graph()) as session:
        input_var = tf.placeholder(
            tf.uint8, (None, ) + image_shape, name=input_name)
        image_var = tf.map_fn(
            lambda x: preprocess_fn(x, is_training=False),
            input_var, back_prop=False, dtype=tf.float32)
        features, _ = network_factory(image_var)
        features = tf.identity(features, name=feature_name)

        saver = tf.train.Saver(slim.get_variables_to_restore())
        saver.restore(session, checkpoint_path)

        output_graph_def = tf.graph_util.convert_variables_to_constants(
            session, tf.get_default_graph().as_graph_def(),
            [features.name.split(":")[0]])
        with tf.gfile.GFile(output_filename, "wb") as file_handle:
            file_handle.write(output_graph_def.SerializeToString())


def encode(preprocess_fn, network_factory, checkpoint_path, images_or_filenames,
           batch_size=32, session=None, image_shape=None):
    """

    Parameters
    ----------
    preprocess_fn : Callable[tf.Tensor] -> tf.Tensor
        A callable that applies preprocessing to a given input image tensor of
        dtype tf.uint8 and returns a floating point representation (tf.float32).
    network_factory : Callable[tf.Tensor] -> (tf.Tensor, tf.Tensor)
        A callable that takes as argument a preprocessed input image of dtype
        tf.float32 and returns the feature representation as well as a logits
        tensors. The logits may be set to None if not required by the loss.
    checkpoint_path : str
        Checkpoint file to load.
    images_or_filenames : List[str] | np.ndarray
        Either a list of filenames or an array of images.
    batch_size : Optional[int]
        Optional batch size; defaults to 32.
    session : Optional[tf.Session]
        Optional TensorFlow session. If None, a new session is created.
    image_shape : Tuple[int, int, int] | NoneType
        Image shape (height, width, channels) or None. If None, `train_x` must
        be an array of images such that the shape can be queries from this
        variable.

    Returns
    -------
    np.ndarray

    """
    if image_shape is None:
        assert type(images_or_filenames) == np.ndarray
        image_shape = images_or_filenames.shape[1:]
    elif type(images_or_filenames) == np.ndarray:
        assert images_or_filenames.shape[1:] == image_shape
    read_from_file = type(images_or_filenames) != np.ndarray

    encoder_fn = _create_encoder(
        preprocess_fn, network_factory, image_shape, batch_size, session,
        checkpoint_path, read_from_file)
    features = encoder_fn(images_or_filenames)
    return features


def _create_encoder(preprocess_fn, network_factory, image_shape, batch_size=32,
                    session=None, checkpoint_path=None, read_from_file=False):
    if read_from_file:
        num_channels = image_shape[-1] if len(image_shape) == 3 else 1
        input_var = tf.placeholder(tf.string, (None, ))
        image_var = tf.map_fn(
            lambda x: tf.image.decode_jpeg(
                tf.read_file(x), channels=num_channels),
            input_var, back_prop=False, dtype=tf.uint8)
        image_var = tf.image.resize_images(image_var, image_shape[:2])
    else:
        input_var = tf.placeholder(tf.uint8, (None, ) + image_shape)
        image_var = input_var

    preprocessed_image_var = tf.map_fn(
        lambda x: preprocess_fn(x, is_training=False),
        image_var, back_prop=False, dtype=tf.float32)

    feature_var, _ = network_factory(preprocessed_image_var)
    feature_dim = feature_var.get_shape().as_list()[-1]

    if session is None:
        session = tf.Session()
    if checkpoint_path is not None:
        tf.train.get_or_create_global_step()
        init_assign_op, init_feed_dict = slim.assign_from_checkpoint(
            checkpoint_path, slim.get_model_variables())
        session.run(init_assign_op, feed_dict=init_feed_dict)

    def encoder(data_x):
        out = np.zeros((len(data_x), feature_dim), np.float32)
        queued_trainer.run_in_batches(
            lambda x: session.run(feature_var, feed_dict=x),
            {input_var: data_x}, out, batch_size)
        return out

    return encoder


def _create_softmax_loss(feature_var, logit_var, label_var):
    del feature_var  # Unused variable
    cross_entropy_var = slim.losses.sparse_softmax_cross_entropy(
        logit_var, tf.cast(label_var, tf.int64))
    tf.summary.scalar("cross_entropy_loss", cross_entropy_var)

    accuracy_var = slim.metrics.accuracy(
        tf.cast(tf.argmax(logit_var, 1), tf.int64), label_var)
    tf.summary.scalar("classification accuracy", accuracy_var)


def _create_magnet_loss(feature_var, logit_var, label_var, monitor_mode=False):
    del logit_var  # Unusued variable
    magnet_loss, _, _ = losses.magnet_loss(feature_var, label_var)
    tf.summary.scalar("magnet_loss", magnet_loss)
    if not monitor_mode:
        slim.losses.add_loss(magnet_loss)


def _create_triplet_loss(feature_var, logit_var, label_var, monitor_mode=False):
    del logit_var  # Unusued variables
    triplet_loss = losses.softmargin_triplet_loss(feature_var, label_var)
    tf.summary.scalar("triplet_loss", triplet_loss)
    if not monitor_mode:
        slim.losses.add_loss(triplet_loss)


def _create_loss(
        feature_var, logit_var, label_var, mode, monitor_magnet=True,
        monitor_triplet=True):
    if mode == "cosine-softmax":
        _create_softmax_loss(feature_var, logit_var, label_var)
    elif mode == "magnet":
        _create_magnet_loss(feature_var, logit_var, label_var)
    elif mode == "triplet":
        _create_triplet_loss(feature_var, logit_var, label_var)
    else:
        raise ValueError("Unknown loss mode: '%s'" % mode)

    if monitor_magnet and mode != "magnet":
        _create_magnet_loss(
            feature_var, logit_var, label_var, monitor_mode=monitor_magnet)
    if monitor_triplet and mode != "triplet":
        _create_triplet_loss(
            feature_var, logit_var, label_var, monitor_mode=True)