Kaggle知识点：内存和显存管理

在kernel竞赛中内存和显存都是受限的，这里就非常建议对内存和显存进行监控了。

当前程序内存占用

使用psutil完成：

import psutil
import os

# MB为单位
psutil.Process(os.getpid()).memory_info().rss/1024/1024

Pandas内存占用

使用os完成：

import os
print(sys.getsizeof(df) / 1024 / 1024, 'MB')

Numpy内存占用

使用os完成：

import os
print(sys.getsizeof(arr) / 1024 / 1024, 'MB')

GPU剩余显存

使用subprocess命令行完成查看，如果有多个GPU，返回的是一个list：

import subprocess as sp
import os

def get_gpu_memory():
    _output_to_list = lambda x: x.decode('ascii').split('\n')[:-1]

    ACCEPTABLE_AVAILABLE_MEMORY = 1024
    COMMAND = "nvidia-smi --query-gpu=memory.free --format=csv"
    memory_free_info = _output_to_list(sp.check_output(COMMAND.split()))[1:]
    memory_free_values = [int(x.split()[0]) for i, x in enumerate(memory_free_info)]
    return memory_free_values

清空内存

使用gc完成删除指定变量：

import gc
del X
gc.collect()

清空显存

使用nuba完成：

from numba import cuda
cuda.select_device(0)
cuda.close()

如果使用torch，也可以使用：

torch.cuda.empty_cache()