diff --git a/paddle/pybind/protobuf.cc b/paddle/pybind/protobuf.cc index bfbe177e8f2a95c8835d32077d43819ef1ad00c5..b86185bf5bf9589ac207141ab6c6eb1e4b8e4696 100644 --- a/paddle/pybind/protobuf.cc +++ b/paddle/pybind/protobuf.cc @@ -17,6 +17,25 @@ limitations under the License. */ namespace paddle { namespace pybind { +template +inline std::vector RepeatedToVector( + const google::protobuf::RepeatedField &repeated_field) { + std::vector ret; + ret.reserve(repeated_field.size()); + std::copy( + repeated_field.begin(), repeated_field.end(), std::back_inserter(ret)); + return ret; +} + +template +inline void VectorToRepeated(const std::vector &vec, + RepeatedField *repeated_field) { + repeated_field->Reserve(vec.size()); + for (auto &elem : vec) { + *repeated_field->Add() = elem; + } +} + void BindProgramDesc(py::module &m) { using namespace paddle::framework; // NOLINT py::class_(m, "ProgramDesc", "") @@ -70,10 +89,7 @@ void BindVarDsec(py::module &m) { [](VarDesc &self, const std::string &name) { self.set_name(name); }) .def("set_shape", [](VarDesc &self, const std::vector &dims) { - LoDTensorDesc *lod_tensor_desc = self.mutable_lod_tensor(); - for (const int64_t &i : dims) { - lod_tensor_desc->add_dims(i); - } + VectorToRepeated(dims, self.mutable_lod_tensor()->mutable_dims()); }) .def("set_data_type", [](VarDesc &self, int type_id) { @@ -82,12 +98,7 @@ void BindVarDsec(py::module &m) { }) .def("shape", [](VarDesc &self) { const LoDTensorDesc &lod_tensor_desc = self.lod_tensor(); - int rank = lod_tensor_desc.dims_size(); - std::vector res(rank); - for (int i = 0; i < rank; ++i) { - res[i] = lod_tensor_desc.dims(i); - } - return res; + return RepeatedToVector(lod_tensor_desc.dims()); }); } diff --git a/paddle/pybind/protobuf.h b/paddle/pybind/protobuf.h index de9a008e25d1e7b076eabea4f83fc71dfca881da..2721c128d1290ee0b1246d877d9e5ea9c4ae24ec 100644 --- a/paddle/pybind/protobuf.h +++ b/paddle/pybind/protobuf.h @@ -27,25 +27,6 @@ namespace py = pybind11; namespace paddle { namespace pybind { -template -inline std::vector RepeatedToVector( - const google::protobuf::RepeatedField& repeated_field) { - std::vector ret; - ret.reserve(repeated_field.size()); - std::copy( - repeated_field.begin(), repeated_field.end(), std::back_inserter(ret)); - return ret; -} - -template -inline void VectorToRepeated(const std::vector& vec, - RepeatedField* repeated_field) { - repeated_field->Reserve(vec.size()); - for (auto& elem : vec) { - *repeated_field->Add() = elem; - } -} - void BindProgramDesc(py::module& m); void BindBlockDesc(py::module& m); void BindVarDsec(py::module& m);