1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
|
## @package layer_model_instantiator
# Module caffe2.python.layer_model_instantiator
from caffe2.python import core, schema
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 [l for l in layers if not include_tags.isdisjoint(l.tags)]
def shrink_output_schema(net, out_schema):
if len(out_schema.field_names()) <= 1:
return out_schema
exists = [net.BlobIsDefined(blob) for blob in out_schema.field_blobs()]
return schema.from_column_list(
[
col_name for ok, col_name in
zip(exists, out_schema.field_names()) if ok
],
[
col_type for ok, col_type in
zip(exists, out_schema.field_types()) if ok
],
[
col_blob for ok, col_blob in
zip(exists, out_schema.field_blobs()) if ok
],
[
col_meta for ok, col_meta in
zip(exists, out_schema.field_metadata()) if ok
]
)
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())
output_schema = shrink_output_schema(
predict_net, model.output_schema.clone()
)
predict_net.set_output_record(output_schema)
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
eval_net.set_input_record(input_schema)
output_schema = shrink_output_schema(
eval_net, model.output_schema + model.metrics_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
train_net.set_input_record(input_schema)
output_schema = shrink_output_schema(
train_net, model.output_schema + model.metrics_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)
model.apply_regularizers_on_loss(train_net, train_init_net)
if not model.has_loss():
return train_init_net, train_net
loss = model.loss
grad_map = train_net.AddGradientOperators(loss.field_blobs())
model.apply_post_grad_net_modifiers(train_net, train_init_net, grad_map,
modify_output_record=True)
model.apply_optimizers(train_net, train_init_net, grad_map)
model.apply_regularizers_after_optimizer(train_net, train_init_net, grad_map)
model.apply_final_net_modifiers(train_net, train_init_net, grad_map,
modify_output_record=True)
return train_init_net, train_net
|
