如何控制显存使用

[johnsonkee]

tf1.x里可以使用
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config)
来控制显存的使用

但是在tf2.0里，session已经被去除，请问如何控制显存的使用呢？

谢谢！

[czy36mengfei]

按需使用显存：
tf.config.gpu.set_per_process_memory_growth(enabled=True)

控制显存占比：
tf.config.gpu.set_per_process_memory_fraction(0.4)

[johnsonkee]

谢谢！