forked from huawei/mindspore2022
initial commit: fix topk index
This commit is contained in:
parent
273d54a4ee
commit
cba7b2ed2d
|
|
@ -26,7 +26,7 @@ __global__ void InTopK(const T *predictions, const int32_t *targets, bool *outpu
|
|||
for (; gt_id < batch_size; gt_id += blockDim.x * gridDim.x) {
|
||||
int32_t target_index = targets[gt_id];
|
||||
T predicted_value = predictions[gt_id * class_id_count + target_index];
|
||||
T top_k_smallest_value = top_k_output[k - 1];
|
||||
T top_k_smallest_value = top_k_output[gt_id * k + k - 1];
|
||||
|
||||
output[gt_id] = predicted_value >= top_k_smallest_value;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -32,9 +32,9 @@ class InTopKNet(nn.Cell):
|
|||
def in_top_k(nptype):
|
||||
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
|
||||
|
||||
predictions = Tensor(np.array([[9, 3, 8, 0, 0, 0, 0, 0, 0],
|
||||
predictions = Tensor(np.array([[4, 1, 2, 0, 0, 0, 0, 0, 0],
|
||||
[7, 9, 9, 0, 0, 0, 0, 0, 0],
|
||||
[9, 9, 9, 0, 0, 0, 0, 0, 0]]).astype(nptype))
|
||||
[3, 3, 3, 0, 0, 0, 0, 0, 0]]).astype(nptype))
|
||||
|
||||
k = 1
|
||||
in_top_k_net = InTopKNet(k)
|
||||
|
|
|
|||
Loading…
Reference in New Issue