torch.nn.DataParallel使用细节

仅使用nn.DataParallel,gpu0和gpu1、gpu0和gpu2、gpu0和gpu3等包含gpu0的组合都是可以的,其余组合不行,报错RuntimeError: module must have its parameters and buffers on device cuda:1 (device_ids[0]) but found one of them on device: cuda:0,因为tensor是默认存在gpu0上的。

解决方法:os.environ["CUDA_VISIBLE_DEVICES"]的使用。

若想使用gpu1和gpu2,程序开头多加一条语句 os.environ["CUDA_VISIBLE_DEVICES"] = "1,2",作用是只允许gpu1和gpu2可用,然后使用Model = nn.DataParallel(Model, device_ids=[0,1]),作用是从可用的两个gpu中搜索第0和第1个位置的gpu。

原文地址:https://www.cnblogs.com/zonechen/p/13395914.html