Add New OP: gumbel_softmax #35506
Add New OP: gumbel_softmax #35506
Conversation
|
Thanks for your contribution! |
| name='x_int32', shape=[2, 3], dtype='int32') | ||
| paddle.nn.functional.gumbel_softmax(x_int32) | ||
|
|
||
| self.assertRaises(TypeError, test_dtype) |
There was a problem hiding this comment.
这里单测缺少了
if __name__ == '__main__':
unittest.main()
| } | ||
|
|
||
| protected: | ||
| framework::OpKernelType GetExpectedKernelType( |
There was a problem hiding this comment.
GetExpectedKernelType这里可以不写，只有一个输入X，默认根据X的特征确定Kernel类型
| "hard", | ||
| "(bool, default false) " | ||
| "if True, the returned samples will be discretized as one-hot vectors, " | ||
| "but will be differentiated as if it is the soft sample in autograd") |
| void InferShape(framework::InferShapeContext* ctx) const override { | ||
| OP_INOUT_CHECK(ctx->HasInput("Out"), "Input", "Out", "gumbel_softmax_grad"); | ||
| OP_INOUT_CHECK(ctx->HasInput(framework::GradVarName("Out")), "Input", | ||
| "Out@grad", "gumbel_softmax_grad"); |
| T min_, max_; | ||
| unsigned int seed_; | ||
| unsigned int offset_ = 0; | ||
| __host__ __device__ UniformGenerator(T min, T max, unsigned int seed) |
There was a problem hiding this comment.
__host__ __device__ -> HOSTDEVICE, 有全局的宏定义
| #include "paddle/fluid/framework/op_registry.h" | ||
| #include "paddle/fluid/framework/operator.h" | ||
| #include "paddle/fluid/memory/memcpy.h" | ||
| #include "paddle/fluid/operators/gumbel_softmax_op.h" |
There was a problem hiding this comment.
| int axis_dim = dX->dims()[axis]; | ||
| // allocate memory on device. | ||
| dX->mutable_data<T>(context.GetPlace()); | ||
| if (dX->numel() == 0) { |
There was a problem hiding this comment.
这里的逻辑和softmax是一样的，所以也同样按照softmax算子的做法添加了对dx的边界处理
| const int size_to_axis = SizeToAxis(axis, dX->dims()); | ||
| const int size_from_axis = SizeFromAxis(axis, dX->dims()); | ||
| Tensor dX_2d, Out_2d, dOut_2d; | ||
| dX_2d.ShareDataWith(*dX).Resize({size_to_axis, size_from_axis}); |
There was a problem hiding this comment.
这里需要将out的数据resize到2维进行处理计算结果，使用shareDataWith可以减少copy次数
| self.check_output_customized(self.verify_output) | ||
|
|
||
| def test_check_grad(self): | ||
| self.check_grad(["X"], "Out", max_relative_error=0.01) |
There was a problem hiding this comment.
这里0.01的精度设置是什么根据？有特殊原因吗？默认是0.005
There was a problem hiding this comment.
由于梯度计算逻辑和softmax算子是一样的，这里精度误差0.01取自softmax算子使用的精度误差
| self.check_grad(["X"], "Out", max_relative_error=0.01) | ||
|
|
||
|
|
||
| class TestGumbelSoftmaxOpGrad(unittest.TestCase): |
There was a problem hiding this comment.
这里测的hard为ture和false条件下，反向梯度是否一致，和前边测的反向内容略有不同
| self.x_shape = [2, 3, 4, 5] | ||
| self.x = np.random.uniform(-1., 1., self.x_shape).astype(np.float32) | ||
| self.count_expected = 24 | ||
| self.place = paddle.CUDAPlace(0) \ |
There was a problem hiding this comment.
with cuda的话应该是CPU和CUDA都测试？没有CUDA的话也应该有CPU测试
There was a problem hiding this comment.
这个下边代码是有的，没有cuda的话是cpu测试
YuanRisheng
dismissed stale reviews from zhhsplendid, lanxianghit, and chenwhql
via
6779ffd
* Add New Op: gumbel_softmax * Add New Op: gumbel_softmax * Add New Op: gumbel_softmax (amend) * add __main__ function in unit test * fix bugs when test in windows ci * update en docs * delete reletive error in unit test * delete relative error in unit test * set hard=True in unit test
PR types
New features
PR changes
OPs
Describe
Add new op named Gumbel Softmax for paddle.
This op is used for sampling from the Gumbel-Softmax distribution and optionally discretizes.
Verify the performance of Gumbel softmax operator locally. The input is random tensor (shape = [200,30]) and repeat the test 100 times. Compare Paddle with Pytorch: the average time-consuming of Paddle is 0.0114 seconds, and the average time-consuming of Pytorch is 0.0123 seconds. The performance of Gumbel softmax operator of Paddle is about 7% better than that of Pytorch.
API:
paddle.nn.functional.gumbel_softmax()Usage example:
Doc
Cn Doc PR: PaddlePaddle/docs/3857