Tensorflow多GPU训练方法

TensorFlow 专栏收录该内容
19 篇文章 4 订阅

在《TensorFlow之目标检测API接口调试(超详细)》文章中,介绍了tensorflow API接口调用方法,详细描述了训练的过程,训练的时候用了单GPU训练,tensorflow是支持多GPU训练的,因此,介绍一下多GPU训练的方法。

很简单,只需要修改两个地方:

第一个参数是train.py文件,找到以下参数

flags.DEFINE_integer('num_clones', 1, 'Number of clones to deploy per worker.')

默认为1,如果自己有多个GPU,可以将参数修改为GPU数量,比如说2、4等;当然也可以在运行train.py的时候以参数的形式传入。

第二个参数是训练配置文件,如rfcn_resnet101_coco.config,根据自己使用的具体网络不同,在train_config中,找到batch_size参数,默认值为1,可修改为大点的参数,理论上大点比较好,并且要求是GPU数量的整数倍,比如设置为8、16、32等,这个参数根据显存来决定,显存越大,可设置的越大。

在batch_size一定的情况下,GPU数量越多,训练速度越快

修改完上述两个参数后,就可以调用多个GPU来训练了。

注意batch_size的设置要合理,如果设置过大,会报如下所示的错误

ResourceExhaustedError (see above for traceback): OOM when allocating tensor with shape[8,2048,36,64] and type float on /job:localhost/replica:0/task:0/device:GPU:1 by allocator GPU_1_bfc

如果说想要指定GPU号,则在train.py中添加以下代码

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,3'

表示调用第1块和第4块显卡训练

  • 1
    点赞
  • 4
    评论
  • 3
    收藏
  • 一键三连
    一键三连
  • 扫一扫,分享海报

相关推荐
©️2020 CSDN 皮肤主题: 游动-白 设计师:白松林 返回首页
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、C币套餐、付费专栏及课程。

余额充值