## @package layer_model_instantiator
# Module caffe2.python.layer_model_instantiator
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from caffe2.python import core
from caffe2.python.layers.layers import InstantiationContext
from caffe2.python.layers.tags import Tags


def _filter_layers(layers, include_tags):
    if include_tags is None:
        return layers
    include_tags = set(include_tags)
    return filter(lambda l: not include_tags.isdisjoint(l.tags), layers)


def generate_predict_net(model, include_tags=None):
    predict_net = core.Net('predict_net')

    for layer in _filter_layers(model.layers, include_tags):
        if Tags.EXCLUDE_FROM_PREDICTION not in layer.tags:
            layer.add_operators(
                predict_net, context=InstantiationContext.PREDICTION)

    predict_net.set_input_record(model.input_feature_schema.clone())
    predict_net.set_output_record(model.output_schema.clone())
    return predict_net


def generate_eval_net(model, include_tags=None):
    eval_net = core.Net('eval_net')

    for layer in _filter_layers(model.layers, include_tags):
        if Tags.EXCLUDE_FROM_EVAL not in layer.tags:
            layer.add_operators(eval_net, context=InstantiationContext.EVAL)

    input_schema = model.input_feature_schema + model.trainer_extra_schema
    output_schema = model.output_schema + model.metrics_schema
    eval_net.set_input_record(input_schema)
    eval_net.set_output_record(output_schema)
    return eval_net


def _generate_training_net_only(model, include_tags=None):
    train_net = core.Net('train_net')
    train_init_net = model.create_init_net('train_init_net')

    for layer in _filter_layers(model.layers, include_tags):
        if Tags.EXCLUDE_FROM_TRAIN not in layer.tags:
            layer.add_operators(train_net, train_init_net)

    input_schema = model.input_feature_schema + model.trainer_extra_schema
    output_schema = model.output_schema + model.metrics_schema
    train_net.set_input_record(input_schema)
    train_net.set_output_record(output_schema)
    return train_init_net, train_net


def generate_training_nets_forward_only(model, include_tags=None):
    train_init_net, train_net = _generate_training_net_only(model, include_tags)
    return train_init_net, train_net


def generate_training_nets(model, include_tags=None):
    train_init_net, train_net = _generate_training_net_only(model, include_tags)

    loss = model.loss
    grad_map = train_net.AddGradientOperators(loss.field_blobs())
    model.apply_optimizers(train_net, train_init_net, grad_map)
    return train_init_net, train_net

@@ -258,16 +258,24 @@
 ort core
+, schema
 %0Afrom ca
@@ -578,16 +578,722 @@
 yers)%0A%0A%0A
+def shrink_output_schema(net, out_schema):%0A    if len(out_schema.field_names()) %3C= 1:%0A        return out_schema%0A    exists = %5Bnet.BlobIsDefined(blob) for blob in out_schema.field_blobs()%5D%0A    return schema.from_column_list(%0A        %5B%0A            col_name for ok, col_name in%0A            zip(exists, out_schema.field_names()) if ok%0A        %5D,%0A        %5B%0A            col_type for ok, col_type in%0A            zip(exists, out_schema.field_types()) if ok%0A        %5D,%0A        %5B%0A            col_blob for ok, col_blob in%0A            zip(exists, out_schema.field_blobs()) if ok%0A        %5D,%0A        %5B%0A            col_meta for ok, col_meta in%0A            zip(exists, out_schema.field_metadata()) if ok%0A        %5D%0A    )%0A%0A%0A
 def gene
@@ -1668,24 +1668,121 @@
 ma.clone())%0A
+    output_schema = shrink_output_schema(%0A        predict_net, model.output_schema.clone()%0A    )%0A
     predict_
@@ -1803,22 +1803,16 @@
 _record(
-model.
 output_s
@@ -1812,32 +1812,24 @@
 utput_schema
-.clone()
 )%0A    return
@@ -2192,35 +2192,60 @@
 xtra_schema%0A    
-out
+eval_net.set_input_record(in
 put_schema = mod
@@ -2234,33 +2234,30 @@
 input_schema
- = model.
+)%0A    
 output_schem
@@ -2258,39 +2258,44 @@
 _schema 
-+ model.metrics
+= shrink_output
 _schema
+(
 %0A
+    
     eval
@@ -2302,38 +2302,57 @@
 _net
-.set_input_record(input
+, model.output_schema + model.metrics
 _schema
+%0A    
 )%0A  
@@ -2826,27 +2826,53 @@
 _schema%0A    
-out
+train_net.set_input_record(in
 put_schema =
@@ -2869,25 +2869,22 @@
 t_schema
- = model.
+)%0A    
 output_s
@@ -2889,39 +2889,44 @@
 _schema 
-+ model.metrics
+= shrink_output
 _schema
+(
 %0A
+    
     trai
@@ -2934,38 +2934,57 @@
 _net
-.set_input_record(input
+, model.output_schema + model.metrics
 _schema
+%0A    
 )%0A  

remove non-existing blobs from output_schema in layer_model_instantiator