本文介绍了如何解决TensorFlow在GPU环境下因显存不足导致的Blas GEMM launch failed错误。提供了两种显存管理方法：1) 动态显存增长，通过tf.config.experimental.set_memory_growth在运行时按需分配；2) 固定显存分配，利用tf.config.experimental.set_virtual_device_configuration预先设定显存使用上限。
在默认情况下，tensorflow为避免显存碎片化，会使用所有显存投入到计算。在这种情况，容易出现Blas GEMM launch failed错误。因此，需要对tensorflow的显存使用进行限制。在tensorflow官网中，给出了两种方法：
第一种方法是随调随用，即在初始的时候只使用少量显存，然后在计算过程中逐步增加显存。主要调用函数是tf.config.experimental.set_memory_growth，代码如下：
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
try:
# Currently, memory growth needs to be the same across GPUs
for
1254
345
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
最低0.47元/天 解锁文章
