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

Merge branch 'ceccocats:master' into master

Browse Source
This commit is contained in:
Adam Cuellar
2021-11-09 16:50:40 -05:00
committed by GitHub
parent 46daa78ca9 d6fb6c6af4
commit 31af5396bb
1 changed files with 9 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernels/postprocessing.cu
+9 -4
View File
@@ -46,11 +46,16 @@ void maxElem_kernel(float *src_begin, float *dst_begin, const int n_classes, con
if (i > size) if (i > size)
return; return;
thrust::device_ptr<float> dPbeg ( &src_begin[i*n_classes] ) ; float max = 0;
thrust::device_ptr<float> dPend = dPbeg + n_classes; int max_idx = 0;
thrust::device_ptr<float> result = thrust::max_element(thrust::device,dPbeg, dPend); for( int j = i*n_classes; j < i*n_classes + n_classes; ++j ){
if( src_begin[j] > max ){
max = src_begin[j];
max_idx = j;
}
}
dst_begin[i] = result - dPbeg; dst_begin[i] = max_idx - i*n_classes;
} }
void maxElem(dnnType *src_begin, dnnType *dst_begin, const int c, const int h, const int w){ void maxElem(dnnType *src_begin, dnnType *dst_begin, const int c, const int h, const int w){