最近在做一个多分类模型时,遇到一个下面bug,明明服务器多块GPU上有空闲GPU,却无法运行模型。
RuntimeError: CUDA out of memory. Tried to allocate 384.00 MiB (GPU 0; 31.75 GiB total capacity; 20.64 GiB already allocated; 265.75 MiB free; 20.75 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
后来发现原因是没有指定GPU,但是有时明明指定了GPU还是无法运行,后来才知道,如果你指定的GPU被占用时,它会报默认0号GPU被占用无法分配运行内存。
# 使用8块GPU中的2,3,4,5号(从0开始)
os.environ['CUDA_VISIBLE_DEVICES'] = '2,3,6,7'
device = torch.device('cuda')
if torch.cuda.device_count() > 1:
print('Let\'s use', torch.cuda.device_count(), 'GPUs!')
model = nn.DataParallel(model, device_ids=[0,1]) # 设置使用的GPU为0和1号
model.to(device)
文章出处登录后可见!
已经登录?立即刷新