如何像PyTorch一样在TensorFlow中实现张量赋值操作

将PyTorch实现的多分类真阳性统计函数迁移到TensorFlow时，直接复用PyTorch写法触发类型错误，无法完成张量构造。

原始可运行PyTorch代码 #

import torch

def true_positive(pred, target, num_classes):
    out = []
    for i in range(num_classes):
        out.append(((pred == i) & (target == i)).sum())
    return torch.tensor(out)

# 测试用例
p = torch.tensor([1]) 
t = torch.tensor([2])
n = torch.tensor([2])
y = true_positive(p,t,n)

报错的TensorFlow版本代码 #

import tensorflow as tf

def true_positive(pred, target, num_classes):
    out = []
    for i in range(num_classes):
        out.append(((pred == i) & (target == i)).sum())
    return torch.tensor(out)

# 测试用例
p = tf.constant([1])
t = tf.constant([2])
n = tf.constant([2])
y = true_positive(p,t,n)

核心报错信息 #

TypeError: only integer scalar arrays can be converted to a scalar index

• 测试代码中传入的类别数参数n是TensorFlow的EagerTensor对象，Python内置range()函数仅接受原生整数标量作为入参，无法直接遍历张量类型值，触发索引类型错误。
• 函数末尾仍调用PyTorch的tensor()构造方法，未替换为TensorFlow对应的张量构造API。

循环逻辑对齐版本 #

完全保留原PyTorch的循环写法，仅做框架适配修改，运行结果和PyTorch版本完全一致：

import tensorflow as tf

def true_positive(pred, target, num_classes):
    out = []
    # 将张量类型的类别数转换为Python原生整数
    if tf.is_tensor(num_classes):
        num_classes = int(num_classes.numpy())
    for i in range(num_classes):
        out.append(((pred == i) & (target == i)).sum())
    # 替换为TensorFlow的张量构造方法
    return tf.convert_to_tensor(out)

# 测试用例
p = tf.constant([1])
t = tf.constant([2])
n = tf.constant([2])
y = true_positive(p, t, n)
print(y) # 输出: tf.Tensor([0 0], shape=(2,), dtype=int32)，和PyTorch运行结果一致

向量化高性能版本（兼容静态图） #

如果需要兼容TensorFlow静态图模式、提升大batch场景下的运行效率，可以用向量化操作替换循环，计算逻辑完全等价：

import tensorflow as tf

def true_positive(pred, target, num_classes):
    pred_onehot = tf.one_hot(pred, depth=num_classes)
    target_onehot = tf.one_hot(target, depth=num_classes)
    return tf.reduce_sum(tf.cast(pred_onehot * target_onehot, tf.int32), axis=0)

# 测试用例
p = tf.constant([1])
t = tf.constant([2])
n = tf.constant([2])
y = true_positive(p, t, n)
print(y) # 输出: tf.Tensor([0 0], shape=(2,), dtype=int32)

内容的提问来源于stack exchange，提问作者Anshuman Sinha