如何修改TensorFlow v1 checkpoint指定变量形状并保存全量模型

问题1：当前实现是否能满足修改变量的需求 #

你修改变量形状、均值填充的逻辑是正确的，但当前代码无法得到完整的新checkpoint，核心问题是当前计算图里只定义了v1、v2两个变量，Saver默认只会保存当前计算图中存在的变量，所以输出文件只有这两个变量的大小。

问题2：保存全量参数的实现方案 #

因为你无法直接加载meta图，可通过读取checkpoint原始变量列表的方式构建全量计算图，步骤如下：

步骤1：读取源checkpoint的所有变量信息 #

用TF原生API无需加载计算图即可获取所有变量的名称、形状：

import tensorflow as tf
import numpy as np

# 源checkpoint路径
source_ckpt = "你的源checkpoint路径"
target_ckpt = "新checkpoint保存路径"

# 读取所有变量信息
reader = tf.train.load_checkpoint(source_ckpt)
var_list = tf.train.list_variables(source_ckpt)

步骤2：构建包含全量变量的计算图 #

遍历所有变量，仅修改目标变量的形状，其余变量保持原始形状定义：

tf.reset_default_graph()
all_vars = {}
# 要修改的变量名
target_vars = {
    "output_weights_agg": [5,768],
    "output_bias_agg": [5]
}

for name, shape in var_list:
    if name in target_vars:
        # 目标变量用修改后的形状定义
        all_vars[name] = tf.get_variable(name, shape=target_vars[name], dtype=tf.float32)
    else:
        # 其余变量保持原始形状定义
        all_vars[name] = tf.get_variable(name, shape=shape, dtype=tf.float32)

步骤3：加载原始参数+修改目标变量值 #

# 先构造只加载非目标变量的Saver，避免形状不匹配报错
restore_var_list = {k:v for k,v in all_vars.items() if k not in target_vars}
restore_saver = tf.train.Saver(var_list=restore_var_list)

# 构造保存全量变量的Saver
save_saver = tf.train.Saver(var_list=all_vars.values())

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    # 加载所有非目标变量的原始值
    restore_saver.restore(sess, source_ckpt)
    
    # 处理第一个目标变量output_weights_agg
    old_v1 = reader.get_tensor("output_weights_agg")
    avg_v1 = np.average(old_v1, axis=0)
    new_v1 = np.vstack((old_v1, avg_v1))
    sess.run(tf.assign(all_vars["output_weights_agg"], new_v1))
    
    # 处理第二个目标变量output_bias_agg
    old_v2 = reader.get_tensor("output_bias_agg")
    avg_v2 = np.average(old_v2, axis=0)
    new_v2 = np.append(old_v2, avg_v2)
    sess.run(tf.assign(all_vars["output_bias_agg"], new_v2))
    
    # 保存全量变量
    save_saver.save(sess, target_ckpt)

TPUEstimator适配方案 #

你可以直接在model_fn的变量初始化阶段修改参数：

1. 正常定义模型，把output_weights_agg、output_bias_agg的形状改为[5,768]和[5]
2. 预训练权重加载完成后，手动获取这两个变量的原始值，拼接均值后重新赋值即可，TPUEstimator保存模型时会自动包含全量参数。

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