Caffe2 - (六)CPU/GPU 模式切换

import sys
sys.path.insert(0, '/path/to/caffe2/build')
from caffe2.python import core, workspace, model_helper
from caffe2.proto import caffe2_pb2, caffe2_legacy_pb2

# -------- CPU/GPU 模式切换 -----
workspace.ResetWorkspace()
# device_opts = core.DeviceOption(caffe2_pb2.CPU, 0) # CPU 模式
device_opts = core.DeviceOption(caffe2_pb2.CUDA, 0) # GPU 模式，及 gpuid


# -------- 加载模型和参数 ------------
INIT_NET = 'init_net.pb'
PREDICT_NET = 'predict_net.pb'

init_def = caffe2_pb2.NetDef()
with open(INIT_NET, 'rb') as f:
    init_def.ParseFromString(f.read())
    init_def.device_option.CopyFrom(device_opts)
    workspace.RunNetOnce(init_def.SerializeToString())

net_def = caffe2_pb2.NetDef()
with open(PREDICT_NET, 'rb') as f:
    net_def.ParseFromString(f.read())
    net_def.device_option.CopyFrom(device_opts)
    workspace.CreateNet(net_def.SerializeToString())

name = net_def.name
output_name = net_def.external_output[-1] # 输出 blob 名
input_name = net_def.external_input[0] # 输入 blob 名

# -------- 送入数据 blob -----------
input_data = np.random.rand(2, 3, 227, 227).astype(np.float32) # NCHW
workspace.FeedBlob(input_name, input_data, device_opts) # device_opts：CPU/GPU 模式的选项

# -------- Forward ----------------
workspace.RunNet(name, 1)

# --------- 读取网络计算结果 --------
results = workspace.FetchBlob(output_name)