注:本文针对单个服务器上多块GPU的使用,不是多服务器多GPU的使用。
当你的电脑中有多块GPU时,可以增加batch_size来加快训练速度等。
涉及的代码
import torch
- 设置可以使用的GPU编号
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,3'
# 使用第0和第3编号的GPU - 判断你的电脑或者pytorch是否能使用GPU来加速模型的训练
torch.cuda.is_available()
-->model = model.cuda()
或者device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-->model = model.to(device)
- 判断你的电脑有几块GPU
torch.cuda.device_count()
- torch.nn.DataParallel启用多GPU并行计算
torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)
说明:
该函数实现了在module级别上的数据并行使用,注意batch size要大于GPU的数量。
参数 :
module:需要多GPU训练的模型
device_ids&