import tensorflow as tf

a = tf.random.shuffle(tf.range(5))

a

tf.sort(a, direction='DESCENDING')

# 返回索引

tf.argsort(a, direction='DESCENDING')

idx = tf.argsort(a, direction='DESCENDING')

tf.gather(a, idx)

idx = tf.argsort(a, direction='DESCENDING')

tf.gather(a, idx)

a = tf.random.uniform([3, 3], maxval=10, dtype=tf.int32)

a

tf.sort(a)

tf.sort(a, direction='DESCENDING')

idx = tf.argsort(a)

idx

# 返回前2个值

res = tf.math.top_k(a, 2)

res

res.values

res.indices

prob = tf.constant([[0.1, 0.2, 0.7], [0.2, 0.7, 0.1]])

target = tf.constant([2, 0])

# 概率最大的索引在最前面

k_b = tf.math.top_k(prob, 3).indices

k_b

k_b = tf.transpose(k_b, [1, 0])

k_b

# 对真实值broadcast，与prod比较

target = tf.broadcast_to(target, [3, 2])

target

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)

    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 / batch_size)

        res.append(acc)

    return res

# 10个样本6类

output = tf.random.normal([10, 6])

# 使得所有样本的概率加起来为1

output = tf.math.softmax(output, axis=1)

# 10个样本对应的标记

target = tf.random.uniform([10], maxval=6, dtype=tf.int32)

print(f'prob: {output.numpy()}')

pred = tf.argmax(output, axis=1)

print(f'pred: {pred.numpy()}')

print(f'label: {target.numpy()}')

acc = accuracy(output, target, topk=(1, 2, 3, 4, 5, 6))

print(f'top-1-6 acc: {acc}')