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)