對Tensorflow限制記憶體配置方式

if len(tf.config.list_physical_devices("GPU")) == 0:
    num_cpu = os.cpu_count()
    max_threads = int(os.getenv("NUM_CPU", 0.5*num_cpu))
    if max_threads <= 0 or max_threads > num_cpu:
        max_threads = int(0.5*num_cpu)
    tf.config.threading.set_inter_op_parallelism_threads(max_threads)
    tf.config.threading.set_intra_op_parallelism_threads(max_threads)

for gpu in tf.config.list_physical_devices("GPU"):
    tf.config.experimental.set_memory_growth(gpu, True)
           

tf版本2.5.0