Place some nodes of the same network on GPU and others on CPU?

Using DeviceScope with the relevant DeviceOption (CPU / GPU) device_type before creating the required Node and it's Blobs

simple example:

from caffe2.python import workspace, model_helper
from caffe2.proto import caffe2_pb2
from caffe2.python import core
import numpy as np

m = model_helper.ModelHelper(name="my first net")
data = np.random.rand(16, 100).astype(np.float32)
gpu_device_id = 1
cpu_device_id = -1
with core.DeviceScope(core.DeviceOption(workspace.GpuDeviceType, gpu_device_id)):
    with core.DeviceScope(core.DeviceOption(caffe2_pb2.CPU, cpu_device_id)):
        # Feed relevant blobs
        workspace.FeedBlob("data", data)
        weight = m.param_init_net.XavierFill([], 'fc_w', shape=[10, 100])
        bias = m.param_init_net.ConstantFill([], 'fc_b', shape=[10, ])
        # Create you cpu Node
        fc_1 = m.net.FC(["data", "fc_w", "fc_b"], "fc1")
    # Create GPU Node
    pred = m.net.Sigmoid(fc_1, "pred")
    softmax, loss = m.net.SoftmaxWithLoss([pred, "label"], ["softmax", "loss"])

print(m.net.Proto())

• Don't forget to do the same before Feed the it's blobs otherwise, you op want be able to reach them.

The output is:

name: "my first net"
op {
name: "my first net"
op {
  input: "data"
  input: "fc_w"
  input: "fc_b"
  output: "fc1"
  name: ""
  type: "FC"
  device_option {
    device_type: 0
    device_id: -1
  }
}
op {
  input: "fc1"
  output: "pred"
  name: ""
  type: "Sigmoid"
  device_option {
    device_type: 1
    device_id: 1
  }
}
op {
  input: "pred"
  input: "label"
  output: "softmax"
  output: "loss"
  name: ""
  type: "SoftmaxWithLoss"
  device_option {
    device_type: 1
    device_id: 1
  }
}
external_input: "data"
external_input: "fc_w"
external_input: "fc_b"
external_input: "label"