Develop xiaobingw (PaddlePaddle#23)

* add support extract loss with IpuOptimizerExtractPass
graphcore · Aug 4, 2021 · 4d3637c · 4d3637c
1 parent 305317c
commit 4d3637c
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 6 deletions.
diff --git a/paddle/fluid/framework/ipu/ipu_backend.h b/paddle/fluid/framework/ipu/ipu_backend.h
@@ -42,9 +42,10 @@ namespace framework {
 using ipu::IpuBuildStrategy;
 
 struct Optimizer {
-  std::string type;
+  std::string type_;
+  std::string loss_;
   // as far as we know, attr is usually float
-  std::map<std::string, float> attrs;
+  std::map<std::string, float> attrs_;
 };
 
 class IpuBackend {
@@ -57,18 +58,20 @@ class IpuBackend {
   void Run(const std::vector<const Tensor *> &inputs,
            const std::vector<Tensor *> &outputs);
 
-  std::string GetOptimizerType() { return optimizer_.type; }
+  std::string GetOptimizerType() { return optimizer_.type_; }
 
-  void SetOptimizerType(const std::string &type) { optimizer_.type = type; }
+  void SetOptimizerType(const std::string &type) { optimizer_.type_ = type; }
 
   const std::map<std::string, float> &GetOptimizerAttr() {
-    return optimizer_.attrs;
+    return optimizer_.attrs_;
   }
 
   void SetOptimizerAttr(const std::string &attr, float value) {
-    optimizer_.attrs[attr] = value;
+    optimizer_.attrs_[attr] = value;
   }
 
+  void SetLoss(const std::string &loss) { optimizer_.loss_ = loss; }
+
   std::vector<int64_t> GetTensorShape(const std::string &var_name) {
     return builder_->getTensorShape(tensors_[var_name]);
   }

diff --git a/paddle/fluid/framework/ir/ipu/optimizer_extract_pass.cc b/paddle/fluid/framework/ir/ipu/optimizer_extract_pass.cc
@@ -54,6 +54,20 @@ void IpuOptimizerExtractPass::ApplyImpl(ir::Graph* graph) const {
           }
         }
       }
+
+      if ((op_role == static_cast<int>(framework::OpRole::kLoss))) {
+        auto outputs = node->Op()->Outputs();
+        PADDLE_ENFORCE_EQ(
+            outputs.size(), 1,
+            platform::errors::InvalidArgument("Can only support one loss key"));
+
+        auto losses_name = outputs.begin()->second;
+        PADDLE_ENFORCE_EQ(losses_name.size(), 1,
+                          platform::errors::InvalidArgument(
+                              "Can only support one loss name"));
+
+        ipu_backend->SetLoss(losses_name[0]);
+      }
     }
   }