TF2 Keras的Conv2DTranspose转TFLite后无法在移动GPU运行的问题

差异原因 #

TF2 Keras的Conv2DTranspose和tf.compat.v1.layers.conv2d_transpose的核心区别在于维度处理逻辑：

• TF2的Keras层默认支持动态输入维度，模型构建时不会固化尺寸计算逻辑，转换为TFLite时，为了适配动态维度的输出尺寸推导，会自动插入Shape、Pack等操作来动态计算输出形状。而移动GPU的TFLite delegate（如NNAPI、GPU Delegate）对这类动态维度操作的支持不完善，导致无法加速。
• tf.compat.v1的conv2d_transpose基于静态图设计，默认要求输入尺寸固定，或在构建阶段就将输出尺寸的计算固化为静态操作，不会生成额外的动态维度相关操作，因此转换后的TFLite模型能被移动GPU正常支持。

解决方法 #

1. 固化输入尺寸并调整转换参数 #

构建TF2 Keras模型时，明确指定输入的固定空间尺寸，避免动态维度：

input_layer = tf.keras.layers.Input(shape=(256, 256, 3))  # 固定高度、宽度、通道数
x = tf.keras.layers.Conv2DTranspose(64, (3,3), strides=(2,2), padding='same')(input_layer)
model = tf.keras.Model(inputs=input_layer, outputs=x)

转换TFLite时，关闭新转换器（部分场景下可避免动态维度处理）：

converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.experimental_new_converter = False
tflite_model = converter.convert()

2. 用静态图模式构建模型 #

通过tf.function的input_signature指定固定输入规格，强制模型使用静态维度计算：

class DeconvModel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.deconv = tf.keras.layers.Conv2DTranspose(64, (3,3), strides=(2,2), padding='same')
    
    @tf.function(input_signature=[tf.TensorSpec(shape=(None, 256, 256, 3), dtype=tf.float32)])
    def call(self, inputs):
        return self.deconv(inputs)

model = DeconvModel()
# 用固定尺寸输入跑一次，固化计算图逻辑
_ = model(tf.random.normal((1, 256, 256, 3)))

再执行TFLite转换，此时输出尺寸的计算会被固化，不会生成Shape/Pack操作。

3. 手动使用底层API实现静态转置卷积 #

直接调用tf.nn.conv2d_transpose，手动计算并指定输出尺寸，完全规避动态推导：

def static_conv2d_transpose(inputs):
    filters = 64
    kernel_size = (3, 3)
    strides = (2, 2)
    padding = 'SAME'
    
    # 手动计算输出尺寸
    batch_size = tf.shape(inputs)[0]
    out_h = inputs.shape[1] * strides[0]
    out_w = inputs.shape[2] * strides[1]
    output_shape = tf.stack([batch_size, out_h, out_w, filters])
    
    # 定义卷积核变量
    kernel = tf.Variable(
        tf.random.normal(shape=kernel_size + (filters, inputs.shape[-1])),
        trainable=True
    )
    return tf.nn.conv2d_transpose(
        inputs, kernel, output_shape=output_shape,
        strides=strides, padding=padding
    )

input_layer = tf.keras.layers.Input(shape=(256, 256, 3))
x = tf.keras.layers.Lambda(static_conv2d_transpose)(input_layer)
model = tf.keras.Model(inputs=input_layer, outputs=x)

这种方式完全控制尺寸计算逻辑，转换后的TFLite模型无额外动态操作，可正常在移动GPU运行。

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