tensorflow pretrained model

最新推荐文章于 2025-05-04 10:00:00 发布

精灵小弟

最新推荐文章于 2025-05-04 10:00:00 发布

阅读量4.3k

点赞数 2

CC 4.0 BY-SA版权

分类专栏： python

本文链接：https://blog.youkuaiyun.com/swfa1/article/details/72636609

python 专栏收录该内容

3 篇文章

订阅专栏

本文介绍了一种在TensorFlow中加载已训练模型的方法，该方法能够自动忽略形状不匹配的层，并提供了一个实用函数optimistic_restore。此函数通过比较变量名称和形状来决定是否加载特定变量。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

tensorflow 简单的加载训练好的模型，如果某一层形状不一样就自动的省略，也可以用reshape。

def optimistic_restore(session, save_file):
    reader = tf.train.NewCheckpointReader(save_file)
    saved_shapes = reader.get_variable_to_shape_map()
    var_names = sorted([(var.name, var.name.split(':')[0]) for var in tf.global_variables()
            if var.name.split(':')[0] in saved_shapes])
    restore_vars = []
    with tf.variable_scope('', reuse=True):
        for var_name, saved_var_name in var_names:
            curr_var = tf.get_variable(saved_var_name)
            var_shape = curr_var.get_shape().as_list()
            if var_shape == saved_shapes[saved_var_name]:
                restore_vars.append(curr_var)
    saver = tf.train.Saver(restore_vars)
    saver.restore(session, save_file)