diff --git a/paddle/operators/recv_op.cc b/paddle/operators/recv_op.cc index f9ed7516826319da422fbb0af4e5c277afa7ae40..ba0d805110d2de0d4a85d2552e1e2af64fc71fd6 100644 --- a/paddle/operators/recv_op.cc +++ b/paddle/operators/recv_op.cc @@ -33,6 +33,8 @@ limitations under the License. */ namespace paddle { namespace operators { +constexpr char kOptimizeBlock[] = "OptimizeBlock"; + void RunServer(std::shared_ptr service) { service->RunSyncUpdate(); VLOG(4) << "RunServer thread end"; @@ -150,14 +152,12 @@ class RecvOp : public framework::OperatorBase { rpc_service_->Reset(); - std::string program_str = Attr("OptimizeProgram"); - framework::proto::ProgramDesc program_desc; - program_desc.ParseFromString(program_str); - framework::ProgramDesc program(program_desc); + auto *block = Attr(kOptimizeBlock); + auto *program = block->Program(); framework::Executor executor(dev_place); // Run sub graph to get optimized tensor try { - executor.Run(program, &recv_scope, 0, /*global_block*/ + executor.Run(*program, &recv_scope, block->ID(), /*global_block*/ false /*create_local_scope*/, false /*create_vars*/); } catch (std::exception &e) { LOG(ERROR) << "run sub program error " << e.what(); @@ -189,8 +189,8 @@ This operator will recv tensor from send_op "IP address to listen on.") .SetDefault("127.0.0.1:6164") .AddCustomChecker([](const std::string &ip) { return !ip.empty(); }); - AddAttr("OptimizeProgram", "type string", - "Serialized ProgramDesc string for recv to run."); + AddAttr( + kOptimizeBlock, "Serialized ProgramDesc string for recv to run."); AddAttr>( "ParamList", "type list of string", "grad->param name mapping to find which param to optimize.") diff --git a/paddle/operators/send_recv_op_test.cc b/paddle/operators/send_recv_op_test.cc index ea091694798475dfd9631910a750405be950c20c..045a0f5434f339bab345d14881ed05450ce6588d 100644 --- a/paddle/operators/send_recv_op_test.cc +++ b/paddle/operators/send_recv_op_test.cc @@ -130,10 +130,7 @@ void StartServerNet(bool is_sparse) { attrs.insert({"endpoint", std::string("127.0.0.1:6174")}); attrs.insert({"ParamList", std::vector({"Out"})}); attrs.insert({"GradList", std::vector({"x1"})}); - std::string program_proto; - PADDLE_ENFORCE(program.Proto()->SerializeToString(&program_proto)); - - attrs.insert({"OptimizeProgram", program_proto}); + attrs.insert({"OptimizeBlock", block}); recv_op = f::OpRegistry::CreateOp("recv", {{"RX", {"x1"}}}, {}, attrs); recv_op->Run(scope, place); } diff --git a/python/paddle/v2/fluid/distribute_transpiler.py b/python/paddle/v2/fluid/distribute_transpiler.py index 06a7b6fb02f5e38e0762d112492854f027fe66ad..e1b7e341d07ce1a533367bd286fc6071fd03b2e6 100644 --- a/python/paddle/v2/fluid/distribute_transpiler.py +++ b/python/paddle/v2/fluid/distribute_transpiler.py @@ -439,7 +439,7 @@ class DistributeTranspiler: }, # grads to recv outputs={}, attrs={ - "OptimizeProgram": optimize_sub_program.desc, + "OptimizeBlock": optimize_sub_program.global_block(), "endpoint": endpoint, "ParamList": [ p.name diff --git a/python/paddle/v2/fluid/distribute_transpiler_simple.py b/python/paddle/v2/fluid/distribute_transpiler_simple.py index bd88f02bde0c6a58138e20db2b07cbd06cd40ba3..56ffb56b1247646903485e5859b60f63df9b97a2 100644 --- a/python/paddle/v2/fluid/distribute_transpiler_simple.py +++ b/python/paddle/v2/fluid/distribute_transpiler_simple.py @@ -243,7 +243,7 @@ class SimpleDistributeTranspiler: self.param_grad_map[endpoint]["grads"]}, # grads to recv outputs={}, attrs={ - "OptimizeProgram": optimize_sub_program.desc, + "OptimizeBlock": optimize_sub_program.global_block(), "endpoint": endpoint, "ParamList": [p.name for p in self.param_grad_map[endpoint]["params"]],