在多GPU上训练时,如何从在单个GPU上训练的模型中加载权重?

时间:2019-10-09 19:57:02

标签: python tensorflow

我想从预先训练的检查点开始进行模型训练。该模型是在单GPU模式下训练的。我现在正在尝试在多个GPU上训练相同的模型。如何从检查点加载权重以在多GPU模式下训练模型?

如果我使用tf.global_variables()恢复变量,它将返回此值。 restore_var = tf.global_variables()

<tf.Variable 'my_scope/res_conv/weights:0' shape=(7, 7, 2, 64) dtype=float32_ref>

而图中的键名称为:

<tf.Variable 'res_conv/weights:0' shape=(7, 7, 2, 64) dtype=float32_ref>

发生这种情况是因为我为Multi-gpu训练定义了一个可变范围。

如何将这些变量分配给多个GPU的新变量?

            with tf.device('/gpu:%d'%i), tf.variable_scope('my_scope', reuse=tf.AUTO_REUSE) as scope:                    
                self.raw_out[i] = model(args.is_training, i)

                all_trainable += tf.trainable_variables()

                gt[i], prediction[i] = self._predict(i, self.raw_out[i])            
                reduced_loss = self._tower_loss(gt[i], prediction[i], args.weight_decay)              

                # calculate gradients on all the batches in one iteration
                gradients = opt.compute_gradients(reduced_loss, all_trainable)

                self.tower_grads[i]=gradients

0 个答案:

没有答案