当前位置: 代码迷 >> 综合 >> python训练时指定GPU
  详细解决方案

python训练时指定GPU

热度:59   发布时间:2024-02-06 20:05:56.0

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

import osos.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,按需修改

 

  相关解决方案