Pytorch 使用GPU指定卡

一. 手动选择

1.运行的主代码内，最前面添加：

import os
os.environ['CUDA_VISIBLE_DEVICES']='4'

 2.指令中添加：

CUDA_VISIBLE_DEVICES=1,2,3 python 文件.py

二. 自动选择

运行的主代码内，最前面添加：

import os
def find_gpus(nums=6):
    os.system('nvidia-smi -q -d Memory |grep -A4 GPU|grep Free >~/.tmp_free_gpus')
    # If there is no ~ in the path, return the path unchanged
    with open(os.path.expanduser('~/.tmp_free_gpus'), 'r') as lines_txt:
        frees = lines_txt.readlines()
        idx_freeMemory_pair = [(idx, int(x.split()[2])) for idx, x in enumerate(frees) ]
    idx_freeMemory_pair.sort(key=lambda my_tuple:my_tuple[1],reverse=True)
    usingGPUs = [str(idx_memory_pair[0]) for idx_memory_pair in idx_freeMemory_pair[:nums] ]
    usingGPUs = ','.join(usingGPUs)
    print('using GPU idx: #', usingGPUs)
    return usingGPUs
os.environ['CUDA_VISIBLE_DEVICES'] = find_gpus(nums=6)