TensorFlow 2中如何为TensorShape添加None动态批次维度

核心原理：你当前得到固定批次为1的张量形状，本质是新增维度时将该维度的静态形状硬编码为1，只要调整新增维度的静态形状规则即可实现动态批次维度。

手动给单样本张量加动态批次维度 #

如果你是要给独立的单样本张量加批次维度，使用以下操作：

import tensorflow as tf

# 原始单样本张量，shape为TensorShape(60, 60, 64)
raw_tensor = tf.random.normal(shape=(60, 60, 64))
# 先新增第一个维度，此时默认静态shape为(1, 60, 60, 64)
batched_tensor = raw_tensor[tf.newaxis, ...]
# 手动覆盖静态形状，将第一维设为动态（None）
batched_tensor = tf.ensure_shape(batched_tensor, (None, 60, 60, 64))

执行后打印batched_tensor.shape即可得到TensorShape([None, 60, 60, 64])。

模型输入场景直接定义动态批次 #

如果你是要将该张量作为模型输入，无需手动加维度，直接在Keras输入层指定单样本形状即可，框架会自动补全动态批次维度：

from tensorflow import keras

# shape参数仅填写单样本的维度，会自动将第一维设为None
input_layer = keras.Input(shape=(60, 60, 64))
print(input_layer.shape) # 输出TensorShape([None, 60, 60, 64])

数据集加载场景的适配 #

如果是tf.data.Dataset加载数据时输出固定批次1的张量，调整批次配置即可：

• 不要调用batch(1)硬写死批次大小
• 如果需要动态适配不同批次大小，可按如下配置：

# 假设原始数据集每个样本的shape为(60, 60, 64)
ds = ds.batch(batch_size=tf.constant(32), drop_remainder=False)
# 此时数据集输出的张量shape为(None, 60, 60, 64)

避坑提示 #

• 不要使用tf.reshape((1, 60, 60, 64))这类写死第一维的操作，会直接把该维度的静态形状固定为1
• 如果是导出模型后出现固定批次1的问题，导出时指定input_signature即可修复：

input_signature = [tf.TensorSpec(shape=(None, 60, 60, 64), dtype=tf.float32)]
tf.saved_model.save(model, "./saved_model", signatures=model.call.get_concrete_function(input_signature))

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