未验证 提交 06694953 编写于 作者: S sneaxiy 提交者: GitHub

Make size op do no data transformation (#42204)

* fix size op

* update
上级 27cb52a4
...@@ -23,6 +23,19 @@ namespace operators { ...@@ -23,6 +23,19 @@ namespace operators {
class SizeOp : public framework::OperatorWithKernel { class SizeOp : public framework::OperatorWithKernel {
public: public:
using framework::OperatorWithKernel::OperatorWithKernel; using framework::OperatorWithKernel::OperatorWithKernel;
protected:
framework::OpKernelType GetExpectedKernelType(
const framework::ExecutionContext& ctx) const override {
auto dtype = framework::proto::VarType::FP32; // dtype is not important
return framework::OpKernelType(dtype, ctx.GetPlace());
}
framework::OpKernelType GetKernelTypeForVar(
const std::string& var_name, const framework::Tensor& tensor,
const framework::OpKernelType& expected_kernel_type) const override {
return expected_kernel_type;
}
}; };
class SizeOpMaker : public framework::OpProtoAndCheckerMaker { class SizeOpMaker : public framework::OpProtoAndCheckerMaker {
...@@ -40,6 +53,8 @@ Return the number of elements in the input. ...@@ -40,6 +53,8 @@ Return the number of elements in the input.
} }
}; };
DECLARE_NO_NEED_BUFFER_VARS_INFERER(SizeOpNoNeedBufferVarInferer, "Input");
} // namespace operators } // namespace operators
} // namespace paddle } // namespace paddle
...@@ -50,4 +65,4 @@ REGISTER_OPERATOR( ...@@ -50,4 +65,4 @@ REGISTER_OPERATOR(
size, ops::SizeOp, ops::SizeOpMaker, size, ops::SizeOp, ops::SizeOpMaker,
paddle::framework::EmptyGradOpMaker<paddle::framework::OpDesc>, paddle::framework::EmptyGradOpMaker<paddle::framework::OpDesc>,
paddle::framework::EmptyGradOpMaker<paddle::imperative::OpBase>, paddle::framework::EmptyGradOpMaker<paddle::imperative::OpBase>,
SizeInferShapeFunctor); SizeInferShapeFunctor, ops::SizeOpNoNeedBufferVarInferer);
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册