使用多GPU同时工作

使用GPU同时工作

在训练神经网络的时候，经常会遇到MemoryError的问题，如果有两块GPU，可以同时使用他们了。

这个程序使用的是下面方法

import argparse

parser.add_argument('--device_ids', default='0,1', type=str, help='comma'
                    ' separated indices of GPU to use, e.g. 0,1 for using GPU_0'
                    ' and GPU_1, default 0,1') # 注意此处 default='0,1' 修改成 0，1代表使用的是0，1两块显卡同时工作


# 主程序中这样设置，即可完成用两个GPU同时训练
model = MODELS[cfg['model']](num_nodes=grid_size, use_crf=cfg['use_crf'])
model = DataParallel(model, device_ids=None)
model = model.cuda()

如果仍然遇到 MemoryError

可以尝试将 Batch_size 改小