TensorFlow 中计算张量中元素个数 tf.size 的基本用法及实例代码

一、环境

TensorFlow API r1.12

CUDA 9.2 V9.2.148

cudnn64_7.dll

Python 3.6.3

Windows 10

 

二、官方说明

返回输入张量的大小，注意不是 shape，而是元素个数，0 维张量，默认为 tf.int32 类型

效果等同于 numpy 中的 np.size()

tf.size(
    input,
    name=None,
    out_type=tf.int32
)

 

参数：

input：张量或稀疏张量

name：可选参数，操作的名称

out_type：可选参数，操作的指定非量化数字输出类型，默认是 tf.int32 类型

返回：

out_type 数据类型的张量，默认是 tf.int32.

 

三、实例

>>> import tensorflow as tf
>>> input_tensor = tf.constant(value=[[[1,1,1],[2,2,2],[3,3,3],[4,4,4]]])
>>> input_tensor
# <tf.Tensor 'Const:0' shape=(1, 4, 3) dtype=int32>
>>> num_elements = tf.size(input_tensor)
>>> num_elements
# <tf.Tensor 'Size:0' shape=() dtype=int32>
>>> with tf.Session() as sess:
...     print(sess.run(num_elements))
...
12