Commit 34e9d1eb authored by Wu Zhiwen's avatar Wu Zhiwen

dnn/Vulkan: support log softmax

Signed-off-by: 's avatarWu Zhiwen <zhiwen.wu@intel.com>
parent 3914c17b
...@@ -18,6 +18,7 @@ layout(push_constant) uniform pushBlock { ...@@ -18,6 +18,7 @@ layout(push_constant) uniform pushBlock {
int channel_size; int channel_size;
int outer_size; int outer_size;
int channels; int channels;
int logsoftmax;
} p; } p;
layout(local_size_x = LOCAL_SZ_X, local_size_y = 1, local_size_z = 1) in; layout(local_size_x = LOCAL_SZ_X, local_size_y = 1, local_size_z = 1) in;
...@@ -68,9 +69,8 @@ void main() ...@@ -68,9 +69,8 @@ void main()
for (int i = 0; i < p.channel_size; ++i) for (int i = 0; i < p.channel_size; ++i)
{ {
float v = output_buffer[index] / sum_buffer[reduced_buffer_off + i]; float v = output_buffer[index] / sum_buffer[reduced_buffer_off + i];
#ifdef LOG_SOFTMAX if (p.logsoftmax == 1)
v = log(v); v = log(v);
#endif
output_buffer[index] = v; output_buffer[index] = v;
index++; index++;
} }
......
...@@ -16,7 +16,7 @@ extern const unsigned int permute_spv[765]; ...@@ -16,7 +16,7 @@ extern const unsigned int permute_spv[765];
extern const unsigned int lrn_spv[1845]; extern const unsigned int lrn_spv[1845];
extern const unsigned int concat_spv[541]; extern const unsigned int concat_spv[541];
extern const unsigned int avg_pool_spv[1538]; extern const unsigned int avg_pool_spv[1538];
extern const unsigned int softmax_spv[1440]; extern const unsigned int softmax_spv[1496];
extern const unsigned int prior_box_spv[1480]; extern const unsigned int prior_box_spv[1480];
extern const unsigned int max_pool_spv[1449]; extern const unsigned int max_pool_spv[1449];
extern const unsigned int relu_spv[502]; extern const unsigned int relu_spv[502];
......
...@@ -22,6 +22,7 @@ struct SoftmaxParam { ...@@ -22,6 +22,7 @@ struct SoftmaxParam {
int channel_size; int channel_size;
int outer_size; int outer_size;
int channels; int channels;
int logsoftmax;
}; };
OpSoftmax::OpSoftmax(const int axis, const bool log_softmax) OpSoftmax::OpSoftmax(const int axis, const bool log_softmax)
...@@ -90,7 +91,7 @@ bool OpSoftmax::forward(Tensor& in, Tensor& out) ...@@ -90,7 +91,7 @@ bool OpSoftmax::forward(Tensor& in, Tensor& out)
bindTensor(device_, *max_tensor_, 1, descriptor_set_); bindTensor(device_, *max_tensor_, 1, descriptor_set_);
bindTensor(device_, *sum_tensor_, 2, descriptor_set_); bindTensor(device_, *sum_tensor_, 2, descriptor_set_);
bindTensor(device_, out, 3, descriptor_set_); bindTensor(device_, out, 3, descriptor_set_);
SoftmaxParam param = {channel_size_, outer_size_, channels_}; SoftmaxParam param = {channel_size_, outer_size_, channels_, log_softmax_ == true ? 1 : 0};
recordCommandBuffer((void *)&param, sizeof(SoftmaxParam)); recordCommandBuffer((void *)&param, sizeof(SoftmaxParam));
runCommandBuffer(); runCommandBuffer();
return true; return true;
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment