提交 3909108c 编写于 作者: Z zhhsplendid

Add SpectralNormGradOpDescMaker

Use SpectralNormGradOpDescMaker instead of DefaultGradOpDescMaker
to avoid registering useless variables to improve GPU usage.

test=develop
上级 85e1cc1e
...@@ -10,6 +10,9 @@ ...@@ -10,6 +10,9 @@
limitations under the License. */ limitations under the License. */
#include "paddle/fluid/operators/spectral_norm_op.h" #include "paddle/fluid/operators/spectral_norm_op.h"
#include <memory>
#include "paddle/fluid/framework/op_registry.h" #include "paddle/fluid/framework/op_registry.h"
namespace paddle { namespace paddle {
...@@ -156,6 +159,28 @@ class SpectralNormOpMaker : public framework::OpProtoAndCheckerMaker { ...@@ -156,6 +159,28 @@ class SpectralNormOpMaker : public framework::OpProtoAndCheckerMaker {
} }
}; };
class SpectralNormGradOpDescMaker : public framework::SingleGradOpDescMaker {
public:
using framework::SingleGradOpDescMaker::SingleGradOpDescMaker;
protected:
std::unique_ptr<framework::OpDesc> Apply() const override {
std::unique_ptr<framework::OpDesc> op(new framework::OpDesc());
op->SetType("spectral_norm_grad");
op->SetInput(framework::GradVarName("Out"), OutputGrad("Out"));
op->SetInput("Weight", Input("Weight"));
op->SetInput("U", Input("U"));
op->SetInput("V", Input("V"));
op->SetOutput(framework::GradVarName("Weight"), InputGrad("Weight"));
op->SetAttrMap(Attrs());
return op;
}
};
class SpectralNormOpGrad : public framework::OperatorWithKernel { class SpectralNormOpGrad : public framework::OperatorWithKernel {
public: public:
using framework::OperatorWithKernel::OperatorWithKernel; using framework::OperatorWithKernel::OperatorWithKernel;
...@@ -185,7 +210,7 @@ class SpectralNormOpGrad : public framework::OperatorWithKernel { ...@@ -185,7 +210,7 @@ class SpectralNormOpGrad : public framework::OperatorWithKernel {
namespace ops = paddle::operators; namespace ops = paddle::operators;
REGISTER_OPERATOR(spectral_norm, ops::SpectralNormOp, ops::SpectralNormOpMaker, REGISTER_OPERATOR(spectral_norm, ops::SpectralNormOp, ops::SpectralNormOpMaker,
paddle::framework::DefaultGradOpDescMaker<true>); ops::SpectralNormGradOpDescMaker);
REGISTER_OPERATOR(spectral_norm_grad, ops::SpectralNormOpGrad); REGISTER_OPERATOR(spectral_norm_grad, ops::SpectralNormOpGrad);
REGISTER_OP_CPU_KERNEL( REGISTER_OP_CPU_KERNEL(
spectral_norm, spectral_norm,
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册