delete lbn2shape setter and getter in register_desc

oneflow/core/graph/boxing_task_node.cpp

@@ -143,9 +143,9 @@ void BoxingTaskNode::FwBuildChainSortedEdgesPair(
     CHECK_EQ(lbns.size(), 1);
     lbns.clear();
     auto in_regst_0 = GetRelatedRegst(sorted_in_edges.at(0));
-    for (const auto& pair : in_regst_0->lbn2shape()) {
-      lbns.push_back(pair.first);
-    }
+    in_regst_0->ForEachLbn([&](const std::string& lbn) {
+      lbns.push_back(lbn);
+    });
   }
   // Enroll Lbn
   auto middle_regst = GetProducedRegstDesc("middle");

oneflow/core/graph/model_save_comp_task_node.cpp

@@ -16,9 +16,9 @@ void MdSaveCompTaskNode::BuildExecAndEnrollLbn2Regsts(TaskGraph* gph) {
     OperatorConf op_conf;
     op_conf.set_name("model_save_op" + updt_task->node_id_str());
     op_conf.mutable_model_save_conf();
-    for (const auto& pair : GetRelatedRegst(SoleInEdge())->lbn2shape()) {
-      op_conf.mutable_model_save_conf()->add_lbns(pair.first);
-    }
+    GetRelatedRegst(SoleInEdge())->ForEachLbn([&](const std::string& lbn) {
+      op_conf.mutable_model_save_conf()->add_lbns(lbn);
+    });
 
     ExecNode* exec_node = mut_exec_gph().NewNode();
     exec_node->mut_op() = OpMgr::Singleton().ConstructOp(op_conf);

oneflow/core/graph/task_node.cpp

@@ -77,7 +77,7 @@ void TaskNode::TakeOverRegstDesc(TaskNode* rhs,
 void TaskNode::EraseProducedEmptyRegsts() {
   EraseIf<std::string, std::shared_ptr<RegstDesc>> (&produced_regst_descs_, []
       (HashMap<std::string, std::shared_ptr<RegstDesc>>::iterator it) {
-    return it->second->lbn2shape().empty();
+    return it->second->NumOfLbn() == 0;
   });
 }
 

oneflow/core/register/register_desc.cpp

@@ -68,13 +68,10 @@ Shape* RegstDesc::GetMutShapePtr(const std::string& lbn) {
   return lbn2shape_.at(lbn).get();
 }
 
-HashMap<std::string, std::unique_ptr<Shape>>& RegstDesc::mut_lbn2shape() {
-  return lbn2shape_;
-}
-
-const HashMap<std::string, std::unique_ptr<Shape>>&
-RegstDesc::lbn2shape() const {
-  return lbn2shape_;
+void RegstDesc::ForEachLbn(std::function<void(const std::string&)> func) const {
+  for (const auto& p : lbn2shape_) {
+    func(p.first);
+  }
 }
 
 void RegstDesc::EraseZeroSizeBlob() {

oneflow/core/register/register_desc.h

@@ -33,8 +33,8 @@ class RegstDesc final {
   void EnrollLbn(const std::string& lbn);
   const Shape& GetShape(const std::string& lbn) const;
   Shape* GetMutShapePtr(const std::string& lbn);
-  HashMap<std::string, std::unique_ptr<Shape>>& mut_lbn2shape();
-  const HashMap<std::string, std::unique_ptr<Shape>>& lbn2shape() const;
+  void ForEachLbn(std::function<void(const std::string&)> func) const;
+  size_t NumOfLbn() const { return lbn2shape_.size(); }
 
   //
   void EraseZeroSizeBlob();