@ -127,13 +127,13 @@ class UnsqueezeOpMaker : public framework::OpProtoAndCheckerMaker {
} ) ;
AddComment ( R " DOC(
Unsqueeze Operator .
Insert single - dimensional entries to the shape of a tensor .
Takes one required argument axes , a list of dimensions that will be inserted .
Dimension indices in axes are as seen in the output tensor .
For example :
Given a tensor such that tensor with shape [ 3 , 4 , 5 ] ,
Insert single - dimensional entries to the shape of a tensor .
Takes one required argument axes , a list of dimensions that will be inserted .
Dimension indices in axes are as seen in the output tensor .
For example :
Given a tensor such that tensor with shape [ 3 , 4 , 5 ] ,
then Unsqueeze ( tensor , axes = [ 0 , 4 ] ) has shape [ 1 , 3 , 4 , 5 , 1 ]
) DOC " );
}
@ -168,6 +168,112 @@ class UnsqueezeGradOp : public framework::OperatorBase {
}
} ;
// FIXME(zcd): unsqueeze2 adds an intermediate output(XShape) based on
// unsqueeze, the XShape is used to carry the shape and lod of X which
// will be used in unsqueeze_grad, in this way, the framework can reuse
// the memory of X immediately the unsqueeze2_op is finished.
// Considering compatibility issues, we could not fix unsqueeze2_op
class Unsqueeze2OpInferShape : public UnsqueezeOpInferShape {
public :
void operator ( ) ( framework : : InferShapeContext * ctx ) const override {
UnsqueezeOpInferShape : : operator ( ) ( ctx ) ;
PADDLE_ENFORCE ( ctx - > HasOutput ( " XShape " ) ,
" Output(XShape) of Unsqueeze operator should not be null. " ) ;
const auto & x_dims = ctx - > GetInputDim ( " X " ) ;
std : : vector < int64_t > xshape_dims ( x_dims . size ( ) + 1 ) ;
xshape_dims [ 0 ] = 0 ;
for ( int i = 0 ; i < x_dims . size ( ) ; + + i ) {
xshape_dims [ i + 1 ] = x_dims [ i ] ;
}
ctx - > SetOutputDim ( " XShape " , framework : : make_ddim ( xshape_dims ) ) ;
ctx - > ShareLoD ( " X " , /*->*/ " XShape " ) ;
}
} ;
class Unsqueeze2OpMaker : public UnsqueezeOpMaker {
public :
void Make ( ) override {
UnsqueezeOpMaker : : Make ( ) ;
AddOutput ( " XShape " ,
" XShape is just used to store the shape and lod of X, which will "
" be used in UnsqueezeGradOp. " )
. AsIntermediate ( ) ;
}
} ;
class Unsqueeze2Op : public framework : : OperatorBase {
public :
using OperatorBase : : OperatorBase ;
private :
void RunImpl ( const framework : : Scope & scope ,
const platform : : Place & place ) const override {
auto & axes = Attr < std : : vector < int > > ( " axes " ) ;
auto x_dims = scope . FindVar ( Input ( " X " ) ) - > Get < framework : : LoDTensor > ( ) . dims ( ) ;
auto out_dims = Unsqueeze2OpInferShape : : GetOutputShape ( axes , x_dims ) ;
framework : : AttributeMap attrs ;
attrs [ " shape " ] = framework : : vectorize2int ( out_dims ) ;
// Invoke Reshape op.
auto reshape_op = framework : : OpRegistry : : CreateOp (
" reshape2 " , { { " X " , { Input ( " X " ) } } , { " Shape " , { } } } ,
{ { " Out " , { Output ( " Out " ) } } , { " XShape " , { Output ( " XShape " ) } } } , attrs ) ;
reshape_op - > Run ( scope , place ) ;
}
} ;
class Unsqueeze2GradOpMaker : public framework : : SingleGradOpDescMaker {
public :
using framework : : SingleGradOpDescMaker : : SingleGradOpDescMaker ;
std : : unique_ptr < framework : : OpDesc > Apply ( ) const override {
auto * grad_op = new framework : : OpDesc ( ) ;
grad_op - > SetType ( " unsqueeze2_grad " ) ;
grad_op - > SetInput ( " XShape " , Output ( " XShape " ) ) ;
grad_op - > SetInput ( framework : : GradVarName ( " Out " ) , OutputGrad ( " Out " ) ) ;
grad_op - > SetOutput ( framework : : GradVarName ( " X " ) , InputGrad ( " X " ) ) ;
grad_op - > SetAttrMap ( Attrs ( ) ) ;
return std : : unique_ptr < framework : : OpDesc > ( grad_op ) ;
}
} ;
class Unsqueeze2GradInferShape : public framework : : InferShapeBase {
public :
void operator ( ) ( framework : : InferShapeContext * context ) const override {
PADDLE_ENFORCE ( context - > HasInput ( " XShape " ) ,
" Input(XShape) shouldn't be null. " ) ;
PADDLE_ENFORCE ( context - > HasInput ( framework : : GradVarName ( " Out " ) ) ,
" Input(Out@GRAD) shouldn't be null. " ) ;
auto xshape_dims = context - > GetInputDim ( " XShape " ) ;
auto x_dims = framework : : slice_ddim ( xshape_dims , 1 , xshape_dims . size ( ) ) ;
context - > SetOutputDim ( framework : : GradVarName ( " X " ) , x_dims ) ;
context - > ShareLoD ( " XShape " , framework : : GradVarName ( " X " ) ) ;
}
} ;
class Unsqueeze2GradOp : public framework : : OperatorBase {
public :
using OperatorBase : : OperatorBase ;
private :
void RunImpl ( const framework : : Scope & scope ,
const platform : : Place & place ) const override {
auto dx_name = Output ( framework : : GradVarName ( " X " ) ) ;
auto dout_name = Input ( framework : : GradVarName ( " Out " ) ) ;
auto xshape_name = Input ( " XShape " ) ;
auto xshape_dims =
scope . FindVar ( xshape_name ) - > Get < framework : : LoDTensor > ( ) . dims ( ) ;
auto x_dims = framework : : slice_ddim ( xshape_dims , 1 , xshape_dims . size ( ) ) ;
framework : : AttributeMap attrs ;
attrs [ " shape " ] = framework : : vectorize2int ( x_dims ) ;
auto reshape_op = framework : : OpRegistry : : CreateOp (
" reshape2 " , { { " X " , { dout_name } } , { " Shape " , { } } } ,
{ { " Out " , { dx_name } } , { " XShape " , { xshape_name } } } , attrs ) ;
reshape_op - > Run ( scope , place ) ;
}
} ;
} // namespace operators
} // namespace paddle
@ -180,3 +286,8 @@ REGISTER_OPERATOR(unsqueeze, ops::UnsqueezeOp, ops::UnsqueezeOpMaker,
paddle : : framework : : DefaultGradOpDescMaker < true > ) ;
REGISTER_OPERATOR ( unsqueeze_grad , ops : : UnsqueezeGradOp ,
ops : : UnsqueezeGradInferShape ) ;
REGISTER_OPERATOR ( unsqueeze2 , ops : : Unsqueeze2Op , ops : : Unsqueeze2OpMaker ,
ops : : Unsqueeze2OpInferShape , ops : : Unsqueeze2GradOpMaker ) ;
REGISTER_OPERATOR ( unsqueeze2_grad , ops : : Unsqueeze2GradOp ,
ops : : Unsqueeze2GradInferShape ) ;
chengduo
7 years ago
committed by
GitHub