[Tensorflow] GPU版で実行時にCUBLAS_STATUS_ALLOC_FAILEDエラーが出たときの対処

環境

Windows 10
Python 3.6.2
Tensorflow 1.6.0(GPU版)

エラー内容

実行時にエラーが表示され、その先頭の方に以下のようなメッセージが表示された場合。

failed to create cublas handle: CUBLAS_STATUS_ALLOC_FAILED

対処

GPUにてメモリ割り当てに失敗している可能性があるため、メモリ割り当ての最大量を指定することで解決することがあります。

alloc.py

import tensorflow as tf

# GPU搭載メモリが4GBだとすると、その0.5倍の2GBまでを割り当てるよう指定する
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.5)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

per_process_gpu_memory_fractionの値を0から1の間の数値で指定します。
この値が高すぎるとメモリ割り当てに失敗することがあります。そのため0.5や0.8など、全メモリではなく、一部のメモリを使用するようにするとエラー発生の頻度を下げられるようです。