| 10 | namespace caffe { |
| 11 | |
| 12 | void InsertSplits(const NetParameter& param, NetParameter* param_split) { |
| 13 | // Initialize by copying from the input NetParameter. |
| 14 | param_split->CopyFrom(param); |
| 15 | param_split->clear_layer(); |
| 16 | map<string, pair<int, int> > blob_name_to_last_top_idx; |
| 17 | map<pair<int, int>, pair<int, int> > bottom_idx_to_source_top_idx; |
| 18 | map<pair<int, int>, int> top_idx_to_bottom_count; |
| 19 | map<pair<int, int>, float> top_idx_to_loss_weight; |
| 20 | map<pair<int, int>, int> top_idx_to_bottom_split_idx; |
| 21 | map<int, string> layer_idx_to_layer_name; |
| 22 | for (int i = 0; i < param.layer_size(); ++i) { |
| 23 | const LayerParameter& layer_param = param.layer(i); |
| 24 | layer_idx_to_layer_name[i] = layer_param.name(); |
| 25 | for (int j = 0; j < layer_param.bottom_size(); ++j) { |
| 26 | const string& blob_name = layer_param.bottom(j); |
| 27 | if (blob_name_to_last_top_idx.find(blob_name) == |
| 28 | blob_name_to_last_top_idx.end()) { |
| 29 | LOG(FATAL) << "Unknown bottom blob '" << blob_name << "' (layer '" |
| 30 | << layer_param.name() << "', bottom index " << j << ")"; |
| 31 | } |
| 32 | const pair<int, int>& bottom_idx = make_pair(i, j); |
| 33 | const pair<int, int>& top_idx = blob_name_to_last_top_idx[blob_name]; |
| 34 | bottom_idx_to_source_top_idx[bottom_idx] = top_idx; |
| 35 | ++top_idx_to_bottom_count[top_idx]; |
| 36 | } |
| 37 | for (int j = 0; j < layer_param.top_size(); ++j) { |
| 38 | const string& blob_name = layer_param.top(j); |
| 39 | blob_name_to_last_top_idx[blob_name] = make_pair(i, j); |
| 40 | } |
| 41 | // A use of a top blob as a loss should be handled similarly to the use of |
| 42 | // a top blob as a bottom blob to another layer. |
| 43 | const int last_loss = |
| 44 | std::min(layer_param.loss_weight_size(), layer_param.top_size()); |
| 45 | for (int j = 0; j < last_loss; ++j) { |
| 46 | const string& blob_name = layer_param.top(j); |
| 47 | const pair<int, int>& top_idx = blob_name_to_last_top_idx[blob_name]; |
| 48 | top_idx_to_loss_weight[top_idx] = layer_param.loss_weight(j); |
| 49 | if (top_idx_to_loss_weight[top_idx]) { |
| 50 | ++top_idx_to_bottom_count[top_idx]; |
| 51 | } |
| 52 | } |
| 53 | } |
| 54 | for (int i = 0; i < param.layer_size(); ++i) { |
| 55 | LayerParameter* layer_param = param_split->add_layer(); |
| 56 | layer_param->CopyFrom(param.layer(i)); |
| 57 | // Replace any shared bottom blobs with split layer outputs. |
| 58 | for (int j = 0; j < layer_param->bottom_size(); ++j) { |
| 59 | const pair<int, int>& top_idx = |
| 60 | bottom_idx_to_source_top_idx[make_pair(i, j)]; |
| 61 | const int split_count = top_idx_to_bottom_count[top_idx]; |
| 62 | if (split_count > 1) { |
| 63 | const string& layer_name = layer_idx_to_layer_name[top_idx.first]; |
| 64 | const string& blob_name = layer_param->bottom(j); |
| 65 | layer_param->set_bottom(j, SplitBlobName(layer_name, |
| 66 | blob_name, top_idx.second, top_idx_to_bottom_split_idx[top_idx]++)); |
| 67 | } |
| 68 | } |
| 69 | // Create split layer for any top blobs used by other layer as bottom |