在训练的脚本开头添加以下代码,指定使用的GPU,然后进行动态分配内存,重新运行就可以。
import os os.environ['CUDA_VISIBLE_DEVICES'] = '0' config = ConfigProto() config.gpu_options.allow_growth = True session = InteractiveSession(config=config)
当然,出现这种问题也存在cudnn版本问题,但是发现这种情况并不多,大多是上面这种。