Pytorch单机多GPU文件内代码备忘

# 对于id=0,1,2...,n块gpu,在程序内只对其中一部分可见,并重新从0排序
os.environ["CUDA_VISIBLE_DEVICES"] = '1,2'

# device设置,默认使用可见gpu中的第一块
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 想要并行使用的可见gpu编号,必须包含device中的id
deviceIDs=[0,1]

# 可占用内存设置(2048*2, rlimit[1])分别为软约束和硬约束，软约束值必须小于硬约束
rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (2048*2, rlimit[1]))

# 实例化模型
model = your_model_name()

# 将模型设置为并行计算
model = nn.DataParallel(model, device_ids = deviceIDs)

# 将模型放入gpu
model = model.to(device)

# 如果模型多gpu并行，则后续如在train中或eval中调用模型需要转换一次类型
if isinstance(faster_rcnn,nn.DataParallel):
        faster_rcnn = faster_rcnn.module