keras使用不同的设备训练模型

Question

我使用经过培训的模型keras。这个型号不适合GPU内存，所以我想把它分成CPU和GPU，也可能分配在不同机器上的GPU上。

1）因此，如果我想在我的CPU上获得所有tf.variable并在GPU上进行计算，我该怎么做呢。因为有些文章表明我们可以将tf.variable放在CPU上并在GPU上进行计算。 （我说的是预先训练好的keras model）。

2）如果我想将预训练模型的图层分成不同机器上的不同GPU。是否正确使用with tf.device然后将图层添加到我的序列模型中。

vgg16_model = keras.applications.vgg16.VGG16()
cnnModel = Sequential()

for layer in vgg16_model.layers[0:13]:
    cnnModel.add(layer)
for layer in vgg16_model.layers[14:16]:
    cnnModel.add(layer)
for layer in vgg16_model.layers[17:21]:
    cnnModel.add(layer)

with tf.device(........):
   cnnModel.add(Dense(2048, name="compress_1"))
   cnnModel.add(Dense(1024, name="compress_2"))
   cnnModel.add(Dense(512, name="compress_3"))

with tf.device(........):
    model = Sequential()
    model.add(keras.layers.TimeDistributed(cnnModel,input_shape=(10,224,224,3),name="CNN_Model"))

with tf.device(........):
    model.add(keras.layers.LSTM(256,name="lstm1",return_sequences=True))
    model.add(keras.layers.LSTM(128,name="lstm2",return_sequences=True))
    model.add(keras.layers.Flatten())
    model.add(keras.layers.Dense(528))
    model.add(keras.layers.Dense(64))
    model.add(keras.layers.Dense(39,activation='sigmoid'))

非常感谢。

更新 我在想是否可以访问模型的张量变量中的属性variable_device并将其自己设置为CPU。

btw我正在使用tf.keras实现。

UPDATE2 我试过一些在张量板上我可以看到：

这是否意味着计算和训练在CPU中，而数据在GPU中？

Answer 1

是的，这是做到这一点的方法。

请注意，因为在CPU和GPU以及其他计算机之间复制变量需要时间。您可以使用单个计算机忽略GPU，或者仅在完成多个批处理后同步计算机之间的变量（Async SGD），从而获得更高的性能。