https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/common_runtime/gpu/gpu_device.cc找设备初始化信息的地方 Found device ,跟进去,可以发现实际是调用了cuMemGetInfo这个函数获取的信息。
那么我们只需要重新写个so,实现这个函数就可以了,这个函数根据我们自己的逻辑设置最大的内存上限,与此同时由于cuMemGetInfo ( size_t free, size_t total ) 这个函数还返回了free的情况,所以我们也需要记录free的值,可以hook掉cuMemAlloc ,申请的时候记录下各个容器的使用情况,也可以利用外围的监控程序获取信息
最终实现的效果,设置环境变量后tensorflow看到的内存值是我们设置的值
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
root@cnn2-1553057347-worker-0:~# export GPU_MEMORY=2147483648 root@cnn2-1553057347-worker-0:~# python Python 3.5.2 (default, Nov 23 2017, 16:37:01) [GCC 5.4.0 20160609] on linux Type "help", "copyright", "credits" or "license"for more information. >>> import tensorflow as tf >>> sess =tf.Session() 2019-03-22 14:57:33.394861: I tensorflow/core/platform/cpu_feature_guard.cc:140] Your CPU supports instructions that this TensorFlow binary was not compiled to use: SSE4.1 SSE4.2 AVX AVX2 FMA hacking gpu memory from GPU_MEMORY env by lovejoy. mem=<2147483648> 2019-03-22 14:57:33.554589: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1344] Found device 0 with properties: name: Tesla K20c major: 3 minor: 5 memoryClockRate(GHz): 0.7055 pciBusID: 0000:02:00.0 totalMemory: 2.00GiB freeMemory: 1.00GiB ##这个卡原来是4G的显存 2019-03-22 14:57:33.554637: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1423] Adding visible gpu devices: 0 2019-03-22 14:57:33.843047: I tensorflow/core/common_runtime/gpu/gpu_device.cc:911] Device interconnect StreamExecutor with strength 1 edge matrix: 2019-03-22 14:57:33.843094: I tensorflow/core/common_runtime/gpu/gpu_device.cc:917] 0 2019-03-22 14:57:33.843106: I tensorflow/core/common_runtime/gpu/gpu_device.cc:930] 0: N