如何修改TensorFlow Checkpoint中第一层张量形状以适配单通道输入？

核心思路 #

你需要将原第一层卷积核的第3个维度（输入通道维度，索引为2）进行求和，把[3, 3, 10, 32]的张量压缩为[3, 3, 1, 32]，这样就能匹配新的单通道输入[200, 200, 1]。其他层的参数可以直接从原检查点加载，不需要修改。

具体实现步骤（适配你使用的tf.train.Saver，TensorFlow 1.x 版本） #

1. 读取原检查点中的第一层卷积核参数
   你可以直接用tf.train.NewCheckpointReader读取原检查点的参数，不用先完整构建旧图结构：

   import tensorflow as tf
   
   # 读取原检查点的参数
   reader = tf.train.NewCheckpointReader('my_checkpoint.ckpt')
   # 替换成你原网络中第一层卷积核的实际变量名，比如'conv1/kernel'
   original_kernel = reader.get_tensor('conv1/kernel')
   # 偏置项直接复用，不需要修改形状
   original_bias = reader.get_tensor('conv1/bias')
   
   # 对输入通道维度（axis=2）求和，保留维度得到[3,3,1,32]的新卷积核
   new_kernel = tf.reduce_sum(original_kernel, axis=2, keepdims=True)
   

2. 定义修改后的新网络结构
   新网络的第一层输入形状改为[200, 200, 1]，卷积核使用上面生成的new_kernel：

   # 定义新的输入占位符
   input_tensor = tf.placeholder(tf.float32, shape=[None, 200, 200, 1], name='new_input')
   
   # 第一层卷积：使用处理后的新卷积核
   with tf.variable_scope('conv1'):
       conv1 = tf.nn.conv2d(input_tensor, new_kernel, strides=[1,1,1,1], padding='SAME')
       conv1 = tf.nn.bias_add(conv1, original_bias)
       conv1 = tf.nn.relu(conv1)
   
   # 后续网络层完全复用原结构，变量名也要和原网络保持一致
   # ... 这里添加你原网络的其他层，比如 conv2、全连接层 fc1 等 ...
   

3. 加载其他层的预训练参数
   因为只有第一层的卷积核被修改了，其他层可以直接从原检查点加载。我们需要构建Saver对象，指定只加载除第一层卷积核之外的变量：

   # 筛选出需要加载的变量：排除第一层的卷积核（我们已经手动设置）
   vars_to_restore = [var for var in tf.global_variables() if 'conv1/kernel' not in var.name]
   
   # 创建Saver并加载参数
   saver = tf.train.Saver(vars_to_restore)
   
   with tf.Session() as sess:
       # 加载除第一层卷积核外的所有预训练参数
       saver.restore(sess, 'my_checkpoint.ckpt')
       # new_kernel是计算得到的张量，已嵌入图中，无需额外初始化
   
       # 现在就可以用修改后的网络进行推理或训练了
       # ...
   

关键注意事项 #

• 变量名匹配：确保你读取的原检查点中第一层卷积核、偏置的变量名，和新网络conv1作用域下的变量名完全一致。如果原网络和新网络的变量名有差异，可以在Saver构造时传入字典手动映射，比如tf.train.Saver({'原网络变量名': 新网络变量})。
• 求和 vs 平均：你要求对第三维度求和，如果你想让权重量级更接近原网络，也可以考虑用tf.reduce_mean求平均，但按你的需求用reduce_sum即可。
• TF2.x 适配：如果使用TensorFlow 2.x，思路完全一致，只是API换成tf.train.Checkpoint和tf.keras的方式，比如用tf.keras.layers.Conv2D并手动设置kernel_initializer为处理后的张量。

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