วิธีที่ค่อนข้างแยกออกจากกันคือการใช้
import tensorflow as tf
from keras import backend as K
num_cores = 4
if GPU:
num_GPU = 1
num_CPU = 1
if CPU:
num_CPU = 1
num_GPU = 0
config = tf.ConfigProto(intra_op_parallelism_threads=num_cores,
inter_op_parallelism_threads=num_cores,
allow_soft_placement=True,
device_count = {'CPU' : num_CPU,
'GPU' : num_GPU}
)
session = tf.Session(config=config)
K.set_session(session)
ในที่นี้booleans
GPU
และCPU
เราระบุว่าเราต้องการรันโค้ดของเราด้วย GPU หรือ CPU หรือไม่โดยกำหนดจำนวน GPU และ CPU อย่างเข้มงวดเพื่อให้เซสชัน Tensorflow ได้รับอนุญาตให้เข้าถึง ตัวแปรnum_GPU
และnum_CPU
กำหนดค่านี้ num_cores
แล้วกำหนดจำนวนของ CPU แกนใช้ได้สำหรับการใช้งานผ่านทางและintra_op_parallelism_threads
inter_op_parallelism_threads
intra_op_parallelism_threads
สั่งตัวแปรจำนวนเธรดการดำเนินการแบบคู่ขนานในโหนดเดียวในกราฟการคำนวณที่ได้รับอนุญาตให้ใช้งาน (ภายใน) ในขณะที่inter_ops_parallelism_threads
ตัวแปรกำหนดจำนวนเธรดที่สามารถเข้าถึงได้สำหรับการดำเนินการแบบขนานระหว่างโหนดของกราฟการคำนวณ (ระหว่าง)
allow_soft_placement
อนุญาตให้รันการดำเนินการบน CPU หากตรงตามเกณฑ์ใด ๆ ต่อไปนี้:
ไม่มีการใช้ GPU สำหรับการทำงาน
ไม่มีอุปกรณ์ GPU ที่รู้จักหรือลงทะเบียน
จำเป็นต้องระบุตำแหน่งร่วมกับอินพุตอื่น ๆ จาก CPU
ทั้งหมดนี้ดำเนินการในคอนสตรัคเตอร์ของคลาสของฉันก่อนการดำเนินการอื่น ๆ และแยกออกจากโมเดลหรือโค้ดอื่น ๆ ที่ฉันใช้
หมายเหตุ: สิ่งนี้ต้องการ tensorflow-gpu
และcuda
/ cudnn
ต้องติดตั้งเนื่องจากมีการกำหนดตัวเลือกให้ใช้ GPU
อ้างถึง:
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
ตามคำตอบด้านล่าง