tensorflow2关于topk的预测值问题

关于topk的预测值问题，最近看龙龙老师的视频，学习了一个小练习，对随机生成的数预测，并且得到top1—top6的预测值

import  tensorflow as tf
import  os

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
tf.random.set_seed(2467)

def accuracy(output, target, topk=(1,)):
   maxk = max(topk)
    batch_size = target.shape[0]

    pred = tf.math.top_k(output, maxk).indices  # 返回位置数组
    pred = tf.transpose(pred, perm=[1, 0])  # 转置
    target_ = tf.broadcast_to(target, pred.shape)  # 得到相同形状的矩阵
    # [10, b]
    correct = tf.equal(pred, target_)  # 比较两个矩阵

    res = []
    for k in topk:
        correct_k = tf.cast(tf.reshape(correct[:k], [-1]), dtype=tf.float32)
        correct_k = tf.reduce_sum(correct_k)
        acc = float(correct_k* (100.0 / batch_size) )
        res.append(acc)

    return res



output = tf.random.normal([10, 6])
output = tf.math.softmax(output, axis=1)
target = tf.random.uniform([10], maxval=6, dtype=tf.int32)
print('prob:', output.numpy())
pred = tf.argmax(output, axis=1)
print('pred:', pred.numpy())
print('label:', target.numpy())

acc = accuracy(output, target, topk=(1,2,3,4,5,6))
print('top-1-6 acc:', acc)

运行结果

D:\pycharm\project\venv\Scripts\python.exe E:/hexiao/cv/深度学习与TensorFlow入门实战-源码和PPT/lesson16-张量排序/topk.py
prob: [[0.2531028  0.21715645 0.16043884 0.13088997 0.04334084 0.19507112]
 [0.05892418 0.04548918 0.00926314 0.14529602 0.66777605 0.07325139]
 [0.09742808 0.08304427 0.07460099 0.04067176 0.626185   0.07806987]
 [0.20478567 0.12294924 0.12010484 0.1375123  0.3641873  0.05046057]
 [0.11872064 0.31072396 0.12530337 0.15528883 0.21325871 0.07670453]
 [0.01519807 0.09672115 0.1460476  0.00934331 0.5649092  0.16778068]
 [0.04199062 0.18141054 0.06647632 0.6006175  0.03198383 0.07752118]
 [0.0922622  0.2346089  0.13022323 0.16295874 0.05362029 0.3263266 ]
 [0.07019574 0.08611771 0.10912605 0.10521299 0.2152082  0.4141393 ]
 [0.01882888 0.2659769  0.19122465 0.24109262 0.14920163 0.13367529]]
pred: [0 4 4 4 1 4 3 5 5 1]
label: [0 2 3 4 2 4 2 3 5 5]
top-1-6 acc: [40.0, 40.0, 50.0, 70.0, 80.0, 100.0]

Process finished with exit code 0