python训练时指定GPU

一、如果是要在py文件里指定,参照如下方式:

import os

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" #可不写
os.environ["CUDA_VISIBLE_DEVICES"] = "0" #0就是指定GPU 0 跑实验,可按需修改成其他GPU


#设置定量的GPU使用量(如果是TensorFlow框架)

import tensorflow as tf
config = tf.ConfigProto() 
config.gpu_options.per_process_gpu_memory_fraction = 0.9 # 占用GPU90%的显存 
session = tf.Session(config=config)

#设置最小的GPU使用量

config = tf.ConfigProto() 
config.gpu_options.allow_growth = True 
session = tf.Session(config=config)

二、如果是在终端指定,参照如下方式:

CUDA_VISIBLE_DEVICES=0 python your_code.py #0表示GPU0,按需修改

猜你喜欢

转载自blog.csdn.net/qq_36663518/article/details/107815195