[EXEC] Specify external memory directly in memory planning (apache#4558)

* [EXEC] Specify external memory directly in memory planning

* Fix warnig

Author: tqchen

NEWS.md

@@ -11,8 +11,6 @@ MXNet Change Log
   - New operators are encouraged to register their property to NNVM op registry attribute
 - Known features removed limitations to be fixed
   - Bulk segment execution not yet added.
-  - The gradient aggregation optimization hack by switching to addto is not yet added,
-    can harm LSTM if it is constructed by unrolling the graph
 
 ## v0.8
 This is the last release before the NNVM refactor.

example/memcost/Makefile

@@ -3,15 +3,15 @@
 
 no_optimization:
 	@echo "Estimating the cost with no optimization..."
-	@MXNET_EXEC_ENABLE_INPLACE=false MXNET_EXEC_MATCH_RANGE=0 python inception_memcost.py
+	@NNVM_EXEC_ENABLE_INPLACE=false NNVM_EXEC_MATCH_RANGE=0 python inception_memcost.py
 
 with_inplace:
 	@echo "Estimating the cost with inplace optimization..."
-	@MXNET_EXEC_ENABLE_INPLACE=true MXNET_EXEC_MATCH_RANGE=0 python inception_memcost.py
+	@NNVM_EXEC_ENABLE_INPLACE=true MXNET_EXEC_MATCH_RANGE=0 python inception_memcost.py
 
 with_sharing:
 	@echo "Estimating the cost with memory sharing ..."
-	@MXNET_EXEC_ENABLE_INPLACE=false python inception_memcost.py
+	@NNVM_EXEC_ENABLE_INPLACE=false python inception_memcost.py
 
 with_both:
 	@echo "Estimating the cost with all optimizations ..."

example/memcost/inception_memcost.py

@@ -36,7 +36,7 @@ def InceptionFactoryB(data, num_3x3red, num_3x3, num_d3x3red, num_d3x3, name):
     cd3x3 = ConvFactory(data=cd3x3r, num_filter=num_d3x3, kernel=(3, 3), pad=(1, 1), stride=(1, 1), name=('%s_double_3x3_0' % name))
     cd3x3 = ConvFactory(data=cd3x3, num_filter=num_d3x3, kernel=(3, 3), pad=(1, 1), stride=(2, 2), name=('%s_double_3x3_1' % name))
     # pool + proj
-    pooling = mx.symbol.Pooling(data=data, kernel=(3, 3), stride=(2, 2), pool_type="max", name=('max_pool_%s_pool' % name))
+    pooling = mx.symbol.Pooling(data=data, kernel=(3, 3), stride=(2, 2), pad=(1, 1), pool_type="max", name=('max_pool_%s_pool' % name))
     # concat
     concat = mx.symbol.Concat(*[c3x3, cd3x3, pooling], name='ch_concat_%s_chconcat' % name)
     return concat

nnvm

@@ -399,7 +399,18 @@ Graph GraphExecutor::InitGraph(nnvm::Symbol symbol,
   // other initializations
   g = nnvm::pass::InferShape(g, arg_shapes, "__shape__");
   g = nnvm::pass::InferType(g, arg_types, "__dtype__");
-  g = nnvm::ApplyPass(g, "PlanMemory");
+
+  {
+    // memory allocator
+    const int kBadStorageID = -1;
+    const int kExternalStorageID = -2;
+    nnvm::StorageVector arg_storage_id(idx.num_node_entries(), kBadStorageID);
+    for (size_t j = num_forward_outputs_; j < idx.outputs().size(); ++j) {
+      arg_storage_id[idx.entry_id(idx.outputs()[j])] = kExternalStorageID;
+    }
+    g.attrs["storage"] = std::make_shared<dmlc::any>(std::move(arg_storage_id));
+    g = nnvm::ApplyPass(g, "PlanMemory");
+  }
   g = DetectInplaceAddTo(g);
   return g;
 }