NewCheckpointReader()函数在Python中读取模型检查点的正确姿势

在Python中，可以使用tf.train.NewCheckpointReader()函数来读取模型检查点。该函数接受一个检查点文件的路径作为参数，并返回一个CheckpointReader对象，该对象可以用于检查和获取模型中的变量。

下面是一个使用NewCheckpointReader()函数读取模型检查点的示例代码：

import tensorflow as tf

# 设置模型检查点文件的路径
checkpoint_path = "/path/to/checkpoint/model.ckpt"

# 创建CheckpointReader对象
reader = tf.train.NewCheckpointReader(checkpoint_path)

# 获取所有变量的名称和形状
var_to_shape_map = reader.get_variable_to_shape_map()
for var_name in var_to_shape_map:
    print("Variable name:", var_name)
    print("Variable shape:", var_to_shape_map[var_name])

# 获取特定变量的值
var_value = reader.get_tensor("variable_name")
print("Variable value:", var_value)

在上面的代码中，我们首先使用NewCheckpointReader()函数创建了一个CheckpointReader对象，并指定了需要读取的模型检查点文件的路径。然后，我们使用get_variable_to_shape_map()方法获取了所有变量的名称和形状，并使用一个循环遍历打印出来。最后，我们使用get_tensor()方法来获取特定变量的值，并将其打印出来。

需要注意的是，tf.train.NewCheckpointReader()函数只能读取模型的权重，而不能读取计算图或其他模型结构。如果你需要加载整个模型（包括结构和权重），可以使用tf.keras.models.load_model()函数或tf.saved_model.load()函数。

此外，还可以使用tf.train.list_variables()函数获取模型检查点文件中的所有变量的名称和形状，而无需使用NewCheckpointReader()函数。示例如下：

import tensorflow as tf

# 设置模型检查点文件的路径
checkpoint_path = "/path/to/checkpoint/model.ckpt"

# 获取所有变量的名称和形状
var_list = tf.train.list_variables(checkpoint_path)
for var_name, var_shape in var_list:
    print("Variable name:", var_name)
    print("Variable shape:", var_shape)

上述代码中，我们直接使用tf.train.list_variables()函数，传入模型检查点文件的路径，返回一个包含所有变量名称和形状的列表。然后，我们使用一个循环遍历列表，并打印出变量的名称和形状。

总结一下，通过使用tf.train.NewCheckpointReader()函数或tf.train.list_variables()函数，可以方便地读取模型检查点文件中的变量信息，进而进行后续的操作，如查看变量的形状，获取特定变量的值等。