TensorFlow模型GPU正常CPU报错：AvgPoolingOp仅支持NHWC格式

基于TensorFlow搭建TactileNet模型，输入数据格式为(#samples,1,19,512)（channels_first/NCHW格式），已通过以下代码设置默认通道位置：

from tensorflow.keras import backend as K
K.set_image_data_format('channels_first')

模型架构完全适配channels_first格式，GPU后端（Google Colab）训练与评估均正常，但切换到CPU后端时，出现如下错误：

InvalidArgumentError: Graph execution error:
...
Node: 'model_4/average_pooling2d_16/AvgPool'
Default AvgPoolingOp only supports NHWC on device type CPU
 [[{{node model_4/average_pooling2d_16/AvgPool}}]] [Op:__inference_test_function_19677]

TensorFlow CPU后端的原生AvgPoolingOp仅支持NHWC（channels_last）格式，对NCHW格式的池化操作无官方实现支持，导致模型在CPU上执行时触发格式不兼容错误。

以下三种方案可解决该问题，根据部署场景选择：

方案1：切换为NHWC格式部署（推荐） #

CPU对NHWC格式支持更完善，是最稳定的部署方案，需调整数据格式与模型参数：

1. 数据格式转换：将输入数据从(#samples,1,19,512)转换为(#samples,19,512,1)
2. 通道格式设置：移除K.set_image_data_format('channels_first')，或改为K.set_image_data_format('channels_last')
3. 模型参数调整：
   • 输入层修改为：input1 = Input(shape=(19, Samples, 1))
   • 所有BatchNormalization的axis参数从1改为3（NHWC格式通道轴为最后一维）
   • 拼接层修改为：concat = concatenate([sub_block1, sub_block2, sub_block4, sub_block3], axis=3)
   • SENet模块简化：由于NHWC格式通道已在最后一维，可去掉Permute层，直接执行multiply操作，示例：

     # 原NCHW格式SENet代码
     # block2       = Permute(dims=(2,3,1))(block2)
     # excitation1  = multiply([block2, excitation1])
     # excitation1  = Permute(dims=(3,1,2))(excitation1)
     # 修改为NHWC格式代码
     excitation1 = multiply([block2, excitation1])
     

方案2：自定义NCHW格式平均池化层 #

通过Lambda层手动实现NCHW格式的平均池化，绕开原生Op的限制：

1. 导入依赖：

   import tensorflow as tf
   from tensorflow.keras.layers import Lambda
   

2. 定义自定义池化函数：

   def nchw_avg_pool(input_tensor, pool_size=(1,2), padding='same'):
       # 针对NCHW格式，对时间维度（最后一维）做平均池化
       if padding == 'same':
           # 计算padding值，保证输出维度匹配
           pad_len = (input_tensor.shape[-1] % pool_size[1]) // 2
           input_tensor = tf.pad(input_tensor, [[0,0], [0,0], [0,0], [pad_len, pad_len]])
       return tf.nn.pool(
           input_tensor,
           window_shape=pool_size,
           pooling_type='AVG',
           strides=pool_size,
           padding='VALID',
           data_format='NCHW'
       )
   

3. 替换模型中所有AveragePooling2D层：

   # 替换示例：原block1中的池化
   # block1 = AveragePooling2D((1, 2))(block1)
   block1 = Lambda(nchw_avg_pool, arguments={'pool_size':(1,2), 'padding':'same'})(block1)
   

方案3：模型转换优化（适合生产部署） #

在GPU上训练完成后，将模型保存为SavedModel格式，通过TensorFlow Lite或TensorRT进行转换，转换时指定输入格式为NHWC，转换工具会自动处理格式兼容问题，生成CPU友好的部署模型：

# 保存训练好的模型
model.save('tactilenet_savedmodel')

# TensorFlow Lite转换示例
import tensorflow as tf
converter = tf.lite.TFLiteConverter.from_saved_model('tactilenet_savedmodel')
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
tflite_model = converter.convert()
with open('tactilenet.tflite', 'wb') as f:
    f.write(tflite_model)

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