如何在tensorflow.contrib.slim中对某些模型进行设备放置

Question

我想在模型中使用经过预训练的resnet_50，并且要为resnet中的所有操作指定GPU：1。但是我不知道该怎么做。谁能帮我吗？

这是我的代码：

slim = tf.contrib.slim
with slim.arg_scope(nets.resnet_v1.resnet_arg_scope()):
    features, _ = nets.resnet_v1.resnet_v1_50(inputs, num_classes=None)

我尝试

slim = tf.contrib.slim
with slim.arg_scope(nets.resnet_v1.resnet_arg_scope()):
    with tf.device("/gpu:1"):
        features, _ = nets.resnet_v1.resnet_v1_50(inputs, num_classes=None)

它不起作用。我找到了一个名为VariableDeviceChooser的函数，但我不知道如何使用它。

请帮助我〜

Answer 1

如果引用的方法不起作用，则还可以使用arg_scope指定不同ResNet操作的设备：

slim = tf.contrib.slim
with slim.arg_scope([slim.conv2d, slim.batch_norm, slim.max_pool2d], device='/GPU:1'):
     with slim.arg_scope(nets.resnet_v1.resnet_arg_scope()):
          features, _ = nets.resnet_v1.resnet_v1_50(inputs, num_classes=None)

此代码段将告诉slim将所有conv，BN和max_pool操作放到GPU：1上，我相信ResNet仅包含这些操作。