@@ -24,6 +24,7 @@ void build_graph(Tensor& input, Tensor& output, bool comments,
2424 ImplType impl1 = parallel ? kTBB : kDefault ;
2525 ImplType impl2 = parallel ? kSTL : kDefault ;
2626 std::vector<std::shared_ptr<Layer>> layers;
27+ std::vector<bool > layerpostop;
2728
2829 std::string json_file = MODEL_PATH_H5 ;
2930 json model_data = read_json (json_file);
@@ -73,12 +74,14 @@ void build_graph(Tensor& input, Tensor& output, bool comments,
7374 1 , pads, 1 , tmp_values, tmp_bias, impl2);
7475 conv_layer->setName (kConvolution );
7576 layers.push_back (conv_layer);
77+ layerpostop.push_back (false );
7678 if (comments) std::cout << " ConvLayer added to layers." << std::endl;
7779 }
7880 if (layer_type.find (" relu" ) != std::string::npos) {
7981 auto ew_layer = std::make_shared<EWLayer>(" relu" );
8082 ew_layer->setName (kElementWise );
8183 layers.push_back (ew_layer);
84+ layerpostop.push_back (true );
8285 if (comments)
8386 std::cout << " Element wise (relu) added to layers" << std::endl;
8487 }
@@ -99,6 +102,7 @@ void build_graph(Tensor& input, Tensor& output, bool comments,
99102 auto fc_layer = std::make_shared<FCLayer>(tensor, tmp_bias);
100103 fc_layer->setName (kFullyConnected );
101104 layers.push_back (fc_layer);
105+ layerpostop.push_back (false );
102106 if (comments) std::cout << " DenseLayer added to layers." << std::endl;
103107 }
104108
@@ -116,6 +120,7 @@ void build_graph(Tensor& input, Tensor& output, bool comments,
116120 auto pool_layer = std::make_shared<PoolingLayer>(shape, pooltype, impl1);
117121 pool_layer->setName (kPooling );
118122 layers.push_back (pool_layer);
123+ layerpostop.push_back (false );
119124 if (comments) std::cout << " PoolingLayer added to layers." << std::endl;
120125 }
121126
@@ -124,13 +129,15 @@ void build_graph(Tensor& input, Tensor& output, bool comments,
124129 std::make_shared<FlattenLayer>(std::vector<size_t >({0 , 3 , 2 , 1 }));
125130 flatten_layer->setName (kFlatten );
126131 layers.push_back (flatten_layer);
132+ layerpostop.push_back (false );
127133 if (comments) std::cout << " FlattenLayer added to layers." << std::endl;
128134 }
129135
130136 if (layer_type.find (" Dropout" ) != std::string::npos) {
131137 auto dropout_layer = std::make_shared<DropOutLayer>(0.0 );
132138 dropout_layer->setName (kDropout );
133139 layers.push_back (dropout_layer);
140+ layerpostop.push_back (false );
134141 if (comments)
135142 std::cout
136143 << " DropOutLayer added to layers with probability 0.4 (turned "
@@ -155,7 +162,12 @@ void build_graph(Tensor& input, Tensor& output, bool comments,
155162 << std::endl;
156163
157164 for (size_t i = 0 ; i < layers.size () - 1 ; ++i) {
158- graph.makeConnection (*layers[i], *layers[i + 1 ]);
165+ if (layerpostop[i]) {
166+ layers[i - 1 ]->postops .layers .push_back (layers[i].get ());
167+ layers[i - 1 ]->postops .count ++;
168+ graph.makeConnection (*layers[i - 1 ], *layers[i + 1 ]);
169+ } else if (!layerpostop[i + 1 ])
170+ graph.makeConnection (*layers[i], *layers[i + 1 ]);
159171 }
160172
161173 graph.setOutput (*layers.back (), output);
0 commit comments