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

最新推荐文章于 2025-05-17 09:04:15 发布

原创最新推荐文章于 2025-05-17 09:04:15 发布 · 5.8k 阅读

5 ·

CC 4.0 BY-SA版权

文章标签：

#Caffe2

Caffe2 同时被 2 个专栏收录

37 篇文章

订阅专栏

Caffe2

37 篇文章

订阅专栏

本文介绍如何在Caffe2中实现CPU与GPU模式的切换，通过修改设备选项来适应不同硬件环境，确保模型顺利运行。文章提供了一段示例代码，详细展示了初始化模型、设置设备选项、加载模型和参数、运行网络等步骤。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

Caffe2 - CPU/GPU 部署模式切换

在尝试把 Caffe 模型转换到 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)