def model_fn(n_frames):

    def _model_fn(features, labels, mode, params):
        input_tensors_list = []
        for i in range(n_frames):
            frame_id = 'frame_{}'.format(i)
            frame_tensor = tf.map_fn(preproc, features[frame_id], tf.float32)
            frame_tensor = tf.expand_dims(frame_tensor, axis=(- 1))
            frame_tensor = tf.transpose(frame_tensor, [0, 1, 2, 4, 3])
            print(frame_tensor)
            input_tensors_list.append(frame_tensor)
        input_tensor_stream = tf.concat(input_tensors_list, axis=3)
        print(input_tensor_stream)
        is_training = (mode == tf.estimator.ModeKeys.TRAIN)
        logits = AquamanNet(input_tensor_stream, is_training, 2)
        total_loss = None
        loss = None
        train_op = None
        eval_metric_ops = {}
        export_outputs = None
        prediction_dict = {'class': tf.argmax(logits, axis=1, name='predictions')}
        if (mode != tf.estimator.ModeKeys.PREDICT):
            reg_loss = tf.losses.get_regularization_loss()
            tf.summary.scalar('Regularization_Loss', tensor=reg_loss)
            loss = tf.losses.softmax_cross_entropy(onehot_labels=labels['target'], logits=logits)
            tf.summary.scalar('XEntropy_LOSS', tensor=loss)
            total_loss = (loss + reg_loss)
            learning_rate = tf.constant(0.0001, name='fixed_learning_rate')
            optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
            vars_to_train = tf.trainable_variables()
            tf.logging.info('Variables to train: {}'.format(vars_to_train))
            if is_training:
                update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
                with tf.control_dependencies(update_ops):
                    train_op = optimizer.minimize(loss=total_loss, global_step=tf.train.get_global_step(), var_list=vars_to_train)
            eval_metric_ops = metrics(logits, labels['target'])
        else:
            export_outputs = {'logits': tf.estimator.export.PredictOutput(outputs=logits)}
        return tf.estimator.EstimatorSpec(mode=mode, predictions=prediction_dict, loss=total_loss, train_op=train_op, eval_metric_ops=eval_metric_ops, export_outputs=export_outputs)
    return _model_fn
def model_fn(n_frames):

    def _model_fn(features, labels, mode, params):
        input_tensors_list = []
        for i in range(n_frames):
            frame_id = 'frame_{}'.format(i)
            frame_tensor = tf.map_fn(preproc, features[frame_id], tf.float32)
            frame_tensor = tf.expand_dims(frame_tensor, axis=(- 1))
            frame_tensor = tf.transpose(frame_tensor, [0, 1, 2, 4, 3])
            print(frame_tensor)
            input_tensors_list.append(frame_tensor)
        input_tensor_stream = tf.concat(input_tensors_list, axis=3)
        print(input_tensor_stream)
        is_training = (mode == tf.estimator.ModeKeys.TRAIN)
        logits = AquamanNet(input_tensor_stream, is_training, 2)
        total_loss = None
        loss = None
        train_op = None
        eval_metric_ops = {}
        export_outputs = None
        prediction_dict = {'class': tf.argmax(logits, axis=1, name='predictions')}
        if (mode != tf.estimator.ModeKeys.PREDICT):
            reg_loss = tf.losses.get_regularization_loss()
            tf.summary.scalar('Regularization_Loss', tensor=reg_loss)
            loss = tf.losses.softmax_cross_entropy(onehot_labels=labels['target'], logits=logits)
            tf.summary.scalar('XEntropy_LOSS', tensor=loss)
            total_loss = (loss + reg_loss)
            learning_rate = tf.constant(0.0001, name='fixed_learning_rate')
            optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
            vars_to_train = tf.trainable_variables()
            tf.logging.info('Variables to train: {}'.format(vars_to_train))
            if is_training:
                update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
                with tf.control_dependencies(update_ops):
                    train_op = optimizer.minimize(loss=total_loss, global_step=tf.train.get_global_step(), var_list=vars_to_train)
            eval_metric_ops = metrics(logits, labels['target'])
        else:
            export_outputs = {'logits': tf.estimator.export.PredictOutput(outputs=logits)}
        return tf.estimator.EstimatorSpec(mode=mode, predictions=prediction_dict, loss=total_loss, train_op=train_op, eval_metric_ops=eval_metric_ops, export_outputs=export_outputs)
    return _model_fn