caffe的python API如何使用多GPU

最新推荐文章于 2026-05-25 10:13:35 发布

原创最新推荐文章于 2026-05-25 10:13:35 发布 · 5.5k 阅读

1 ·

本内容遵循CC 4.0 BY-SA版权协议

标签

#multiple gpus #caffe python API

收录于

本文介绍了如何在Caffe框架中配置多GPU进行模型训练。由于Caffe本身并未提供直接支持多GPU训练的Python API，文章给出了通过设置特定参数并在配置文件中指定GPU设备的方式来实现多GPU并行训练的方法。

准备跑clone下来的faster rcnn时，本打算使用多个gpu进行训练，以为--gpu 1，2，3就可以了，too naive！

caffe的官方网站给出的说明就是：

caffe.set_mode_gpu()
caffe.set_device(0)

更多相关信息可以参考这个链接： https://www.programcreek.com/python/example/83262/caffe.set_mode_gpu

实际上也确实没有pythonAPI调用多个gpu的函数。这点可以参考caffe的官方网站的讨论话题。

如何再python下调用多个gpu呢，SSD给了一个范例：

from caffe.model_libs import *  
# Use GPU or CPU
solver_mode = P.Solver.GPU
# Defining which GPUs to use.
gpus = "0"

再然后，将这些所有的参数写入到配置文件中去，在c++中调用多gpu

# Create job file.
with open(job_file, 'w') as f:
  f.write('cd {}\n'.format(caffe_root))
  f.write('./build/tools/caffe test \\\n')
  f.write('--model="{}" \\\n'.format(test_net_file))
  f.write('--weights="{}" \\\n'.format(pretrain_model))
  f.write('--iterations="{}" \\\n'.format(test_iter))
  if solver_mode == P.Solver.GPU:
    f.write('--gpu {} 2>&1 | tee {}/{}.log\n'.format(gpus, job_dir, model_name))