tensorflow、pytorch命令如何指定GPU的使用

一、tensorflow如何使用GPU

1、指定使用某块GPU

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0' #指定第一块GPU可用
os.environ[‘CUDA_VISIBLE_DEVICES’] = '0, 1' #指定第一、二块GPU可用

2、分配GPU的使用情况

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.5 #程序最多只能占用指定gpu50%的显存
config.gpu_options.allow_growth = True #程序按需申请内存
sess = tf.Session(config = config)

二、pytorch中指定多块GPU同时使用

   opt.device = t.device('cuda') if opt.use_gpu else t.device('cpu')
   net = t.nn.DataParallel(model, device_ids=[0, 1])
   net.to(opt.device)