mmr/tkDNN
Archived
1
0
Fork 0
Code Issues 46 Pull Requests 12 Releases 4 Activity

Post-processing improved

Browse Source 
Signed-off-by: Micaela Verucchi <micaelaverucchi@gmail.com>
This commit is contained in:
Micaela Verucchi
2020-06-24 18:46:25 +02:00
parent 082920f3f5
commit d25803d438
4 changed files with 52 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernels/postprocessing.cu
+19
View File
@@ -34,6 +34,25 @@ void sortAndTopKonDevice(dnnType *src_begin, int *idsrc, float *topk_scores, int
sortAndTopK_kernel<<<blocks, threads, 0>>>(src_begin, idsrc, topk_scores, topk_inds, topk_ys, topk_xs, size, K);
}
__global__
void maxElem_kernel(float *src_begin, float *dst_begin, const int n_classes, const int size){
int i = blockDim.x*blockIdx.x + threadIdx.x;
if (i > size)
return;
thrust::device_ptr<float> dPbeg ( &src_begin[i*n_classes] ) ;
thrust::device_ptr<float> dPend = dPbeg + n_classes;
thrust::device_ptr<float> result = thrust::max_element(thrust::device,dPbeg, dPend);
dst_begin[i] = result - dPbeg;
}
void maxElem(dnnType *src_begin, dnnType *dst_begin, const int c, const int h, const int w){
int blocks = (h*w)/32+1;
int threads = 32;
maxElem_kernel<<<blocks, threads, 0>>>(src_begin, dst_begin, c, h*w);
}
void topKxyclasses(int *ids_begin, int *ids_end, const int K, const int size, const int wh, int *clses, int *xs, int *ys){
thrust::transform(thrust::device, ids_begin, ids_end, thrust::make_constant_iterator(wh), clses, thrust::divides<int>());
thrust::transform(thrust::device, ids_begin, ids_end, thrust::make_constant_iterator(wh), ids_begin, thrust::modulus<int>());