如何在Tensorflow 2.0中加载模型权重以进行多次gpu训练?

时间:2020-06-25 23:40:06

标签: python tensorflow

我正在尝试在TensorFlow 2.0中加载模型权重以进行多次GPU训练。我的模型检查点来自Tensorflow 1.0中的训练。

这就是我以前在TF1中加载权重的方式。

var_list = tf.trainable_variables()
for var in var_list:           
    vname = ''.join(var.name.split('semantic/'))
    val = tf.contrib.framework.load_variable(checkpoint, vname)
    tf.assign(var, val)

我需要类似的方法来在tf2中加载权重。感谢您的帮助。

谢谢!

1 个答案:

答案 0 :(得分:0)

尝试:

typedef enum {
  HW_CONFIG_A,
  HW_CONFIG_B,
  /* ... */
} HWConfig_e;

void myFunc( HWConfig_e cfg, union __my_container *cont )
{
    switch( cfg )
    {
        case HW_CONFIG_A:
        {
             __my_struct_configA *cfgA = &(cont->a);

             printf("a=%d b=%l c=%c\n", cfgA->a, cfgA->b, cfgA->c );
             break;

        case HW_CONFIG_B:
        {
             __my_struct_configB *cfgB = &(cont->b);

             printf("a=%d (b=N/A) c=%c\n", cfgB->a, cfgB->c );
             break;
        }
    }
}


这只是您在tf2中的tf1代码,您可能需要使用类似的

var_list = model.trainable_variables

for var in var_list:      
   vname = ''.join(var.name.split('semantic/'))
   val = tf.train.load_variable(checkpoint, vname)
   var = val

取决于您建立模型的方式

相关问题