如何使用tf.assert判断TensorFlow张量中所有数值均为0或1？

要实现这个需求，我们可以分两步构造断言逻辑：先判断张量里的每个元素是否是0或1，再验证所有元素都满足这个条件，最后用tf.assert_true触发断言检查。

下面结合你给出的示例张量，附上完整实现代码：

import tensorflow as tf

# 你提供的示例张量
bad_mask = tf.Variable([[0.0,1.0,0.2,0.0,0.0], [0.0,5.0,0.0,2.3,0.0]])
good_mask = tf.Variable([[0.0,1.0,1.0,0.0,0.0], [0.0,1.0,0.0,1.0,0.0]])

def assert_binary_tensor(tensor):
    # 生成布尔掩码：标记每个元素是否是0或1
    is_zero_or_one = tf.logical_or(tf.equal(tensor, 0.0), tf.equal(tensor, 1.0))
    # 检查所有元素是否都符合条件
    all_valid = tf.reduce_all(is_zero_or_one)
    # 触发断言，不满足条件时抛出错误
    return tf.assert_true(all_valid, message="张量包含非0且非1的数值！")

# 测试合法张量（不会触发错误）
with tf.control_dependencies([assert_binary_tensor(good_mask)]):
    good_result = tf.identity(good_mask)

# 测试非法张量（会触发断言错误）
with tf.control_dependencies([assert_binary_tensor(bad_mask)]):
    bad_result = tf.identity(bad_mask)

# 在tf.function中运行测试（TF2.x推荐方式）
@tf.function
def test_tensor_validity(tensor):
    assert_binary_tensor(tensor)
    return tensor

# 验证合法张量
print("测试合法张量：")
test_tensor_validity(good_mask)
print("合法张量通过检查！")

# 验证非法张量（取消注释会触发InvalidArgumentError）
# print("\n测试非法张量：")
# test_tensor_validity(bad_mask)

关键细节说明： #

• tf.logical_or(tf.equal(tensor, 0.0), tf.equal(tensor, 1.0))：生成与输入张量同形状的布尔张量，每个位置标记对应元素是否为0或1。
• tf.reduce_all：把布尔张量压缩成一个标量，只有当所有元素都是True时才返回True。
• tf.assert_true：当传入的条件为False时，会抛出InvalidArgumentError，并附带你自定义的错误提示。
• 若你的张量是整数类型，只需把代码里的0.0和1.0换成0和1即可，逻辑完全通用。

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