initial commit: fix topk index

This commit is contained in:
Peilin Wang 2021-07-12 16:05:14 -04:00
parent 273d54a4ee
commit cba7b2ed2d
2 changed files with 3 additions and 3 deletions

View File

@ -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;
}

View File

@ -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)