//  [python]
def model_fn(features, labels, mode, params, model):
    """Initializes a model for the estimator.

    Args:
        features: Input tensor representing a batch of input features
        labels: Labels tensor representing a batch of labels
        mode: Indicates which mode is running
        params: The parameters
        model: The model
    """
    is_training = mode == tf.estimator.ModeKeys.TRAIN
    is_predicting = mode == tf.estimator.ModeKeys.PREDICT

    # Generate a summary node for the images
    if is_training:
        add_image_summaries(features, params)

    # Apply the model, if fp16 cast to fp32 for numerical stability.
    logits, fc7 = model(features, is_training)
    logits = tf.cast(logits, tf.float32)
    fc7 = tf.cast(fc7, tf.float32)

    # If necessary, slice last sequence step for logits
    final_logits = logits[:,-1,:] if logits.get_shape().ndims == 3 else logits
    final_fc7 = fc7[:,-1,:] if fc7.get_shape().ndims == 3 else fc7

    # Derive preds and probs from final_logits
    predictions = {
        'fc7': final_fc7,
        'classes': tf.argmax(final_logits, axis=-1),
        'probabilities': tf.nn.softmax(final_logits, name='softmax_tensor')}

    if is_predicting:
        # Return the predictions and the specification for serving a SavedModel
        return tf.estimator.EstimatorSpec(
            mode=mode,
            predictions=predictions,
            export_outputs={
                'predict': tf.estimator.export.PredictOutput(predictions)
            })

    # If necessary, slice last sequence step for labels
    final_labels = labels[:,-1] if labels.get_shape().ndims == 2 else labels

    def _compute_balanced_sample_weight(labels):
        """Calculate the balanced sample weight for imbalanced data."""
        f_labels = tf.reshape(labels,[-1]) if labels.get_shape().ndims == 2 else labels
        y, idx, count = tf.unique_with_counts(f_labels)
        total_count = tf.size(f_labels)
        label_count = tf.size(y)
        calc_weight = lambda x: tf.divide(tf.divide(total_count, x),
            tf.cast(label_count, tf.float64))
        class_weights = tf.map_fn(fn=calc_weight, elems=count, dtype=tf.float64)
        sample_weights = tf.gather(class_weights, idx)
        sample_weights = tf.reshape(sample_weights, tf.shape(labels))
        return tf.cast(sample_weights, tf.float32)

    # Training with multiple labels per sequence
    if params.use_sequence_loss:

        if is_training:
            sample_weights = _compute_balanced_sample_weight(labels)
        else:
            sample_weights = tf.ones_like(labels, dtype=tf.float32)

        # Calculate and scale cross entropy
        scaled_loss = tf.contrib.seq2seq.sequence_loss(
            logits=logits,
            targets=tf.cast(labels, tf.int32),
            weights=sample_weights)
        tf.identity(scaled_loss, name='seq2seq_loss')
        tf.compat.v1.summary.scalar('loss/seq2seq_loss', scaled_loss)

    # Training on single for frames or with one label per sequence
    else:

        if is_training:
            sample_weights = _compute_balanced_sample_weight(final_labels)
        else:
            sample_weights = tf.ones_like(final_labels, dtype=tf.float32)

        # Calculate scaled cross entropy
        unscaled_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
            labels=tf.cast(final_labels, tf.int32),
            logits=final_logits)
        scaled_loss = tf.reduce_mean(tf.multiply(unscaled_loss, sample_weights))
        tf.compat.v1.summary.scalar('loss/scaled_loss', scaled_loss)

    # Compute loss with Weight decay
    l2_loss = params.l2_lambda * tf.add_n(
        [tf.nn.l2_loss(v) for v in tf.compat.v1.trainable_variables()
            if 'norm' not in v.name])
    tf.compat.v1.summary.scalar('loss/l2_loss', l2_loss)
    loss = scaled_loss + l2_loss

    if is_training:
        global_step = tf.compat.v1.train.get_or_create_global_step()

        def _decay_fn(learning_rate, global_step):
            return tf.compat.v1.train.exponential_decay(
                learning_rate=learning_rate, global_step=global_step,
                decay_steps=params.steps_per_epoch, decay_rate=params.decay_rate)

        def _grad_filter(grad_vars):
            """Only apply gradient updates to the certain layers."""
            if 'dense_lstm' in params.finetune_only:
                return [g_v for g_v in grad_vars \
                    if 'dense' in g_v[1].name or 'lstm' in g_v[1].name]
            elif 'dense' in params.finetune_only:
                return [g_v for g_v in grad_vars if 'dense' in g_v[1].name]
            else:
                return grad_vars

        # Learning rate
        learning_rate = _decay_fn(params.base_learning_rate, global_step)
        tf.identity(learning_rate, name='learning_rate')
        tf.compat.v1.summary.scalar('training/learning_rate', learning_rate)

        # The optimizer
        optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=learning_rate)

        if params.dtype == tf.float16:
            loss_scale_manager = tf.contrib.mixed_precision.ExponentialUpdateLossScaleManager(
                init_loss_scale=params.loss_scale, incr_every_n_steps=5)
            loss_scale_optimizer = tf.contrib.mixed_precision.LossScaleOptimizer(
                optimizer, loss_scale_manager)
            grad_vars = loss_scale_optimizer.compute_gradients(loss)
        else:
            grad_vars = optimizer.compute_gradients(loss)

        # Filter vars to retain fine tuned vars
        grad_vars = _grad_filter(grad_vars)

        tf.compat.v1.summary.scalar("training/global_gradient_norm",
            tf.linalg.global_norm(list(zip(*grad_vars))[0]))

        # Clip gradients
        grads, vars = zip(*grad_vars)
        grads, _ = tf.clip_by_global_norm(grads, params.gradient_clipping_norm)
        grad_vars = list(zip(grads, vars))

        for grad, var in grad_vars:
            var_name = var.name.replace(":", "_")
            tf.compat.v1.summary.histogram("gradients/%s" % var_name, grad)
            tf.compat.v1.summary.scalar("gradient_norm/%s" % var_name, tf.linalg.global_norm([grad]))
        tf.compat.v1.summary.scalar("loss", loss)
        tf.compat.v1.summary.scalar("training/clipped_global_gradient_norm",
            tf.linalg.global_norm(list(zip(*grad_vars))[0]))

        if params.dtype == tf.float16:
            minimize_op = loss_scale_optimizer.apply_gradients(grad_vars, global_step)
        else:
            minimize_op = optimizer.apply_gradients(grad_vars, global_step)

        update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS)
        train_op = tf.group(minimize_op, update_ops)

    else:
        train_op = None

    # Calculate accuracy metrics - always done with final labels
    final_labels = tf.cast(final_labels, tf.int64)
    accuracy = tf.compat.v1.metrics.accuracy(
        labels=final_labels, predictions=predictions['classes'])
    mean_per_class_accuracy = tf.compat.v1.metrics.mean_per_class_accuracy(
        labels=final_labels, predictions=predictions['classes'],
        num_classes=params.num_classes)
    tf.compat.v1.summary.scalar('metrics/accuracy', accuracy[1])
    tf.compat.v1.summary.scalar('metrics/mean_per_class_accuracy',
        tf.reduce_mean(mean_per_class_accuracy[1]))
    metrics = {
        'metrics/accuracy': accuracy,
        'metrics/mean_per_class_accuracy': mean_per_class_accuracy}

    # Calculate class-specific metrics
    for i in range(params.num_classes):
        class_precision = tf.compat.v1.metrics.precision_at_k(
            labels=final_labels, predictions=final_logits, k=1, class_id=i)
        class_recall = tf.metrics.recall_at_k(
            labels=final_labels, predictions=final_logits, k=1, class_id=i)
        tf.compat.v1.summary.scalar('metrics/class_%d_precision' % i, class_precision[1])
        tf.compat.v1.summary.scalar('metrics/class_%d_recall' % i, class_recall[1])
        metrics['metrics/class_%d_precision' % i] = class_precision
        metrics['metrics/class_%d_recall' % i] = class_recall

    # Log number of trainable model params
    trainable_params = [tf.reduce_prod(v.shape) for v in tf.compat.v1.trainable_variables()]
    tf.compat.v1.summary.scalar('model/trainable_params', sum(trainable_params))

    # Log number of total model params
    total_params = [tf.reduce_prod(v.shape) for v in tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.GLOBAL_VARIABLES)]
    tf.compat.v1.summary.scalar('model/total_params', sum(total_params))

    # Log estimated checkpoint size
    weights_size = [tf.reduce_prod(v.shape) * v.dtype.size
        for v in tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.GLOBAL_VARIABLES)]
    tf.compat.v1.summary.scalar('model/weights_size_mb', sum(weights_size) / (1024 ** 2))

    # Log estimated back and forward activation size
    activations_size = [tf.reduce_prod(v.shape) * v.dtype.size * tf.constant(2) \
        * params.batch_size for v in tf.compat.v1.trainable_variables()]
    tf.compat.v1.summary.scalar('model/activations_size_mb', sum(activations_size) / (1024 ** 2))

    return tf.estimator.EstimatorSpec(
        mode=mode,
        predictions=predictions,
        loss=loss,
        train_op=train_op,
        eval_metric_ops=metrics)