pytorch gpu 显存 内存 管理 调试 监控 日志

pytorch gpu 显存调试

如何运行 gpu_memory_log

import torch
from gpu_memory_log import gpu_memory_log

dtype = torch.float
N, D_in, H, D_out = 64, 1000, 100, 10

device = torch.device("cuda")
x = torch.randn(N, D_in, device=device, dtype=dtype)
y = torch.randn(N, D_out, device=device, dtype=dtype)

w1 = torch.randn(D_in, H, device=device, dtype=dtype, requires_grad=True)
w2 = torch.randn(H, D_out, device=device, dtype=dtype, requires_grad=True)

learning_rate = 1e-6
gpu_memory_log()
for t in range(500):
    y_pred = x.mm(w1).clamp(min=0).mm(w2)
    loss = (y_pred - y).pow(2).sum()
    #print(t, loss.item())
    loss.backward()
    with torch.no_grad():
        w1 -= learning_rate * w1.grad
        w2 -= learning_rate * w2.grad

        w1.grad.zero_()
        w2.grad.zero_()

print(loss.item())
gpu_memory_log()

监控pytorch gpu 内存的代码，很简单，就是 gpu_memory_log()
把这行代码放在，想要看gpu显存状态的位置。

output

LINE:27, FUNC:<module>, FILE:gpu_test.py, TIME:2019-06-02 17:16:23.288399, CONTENT:
[tensor: 1 * Size:(100, 10) | Memory: 0.0038 M | <class 'torch.Tensor'>]
[tensor: 1 * Size:(64, 1000) | Memory: 0.2441 M | <class 'torch.Tensor'>]
[tensor: 1 * Size:(64, 10) | Memory: 0.0024 M | <class 'torch.Tensor'>]
[tensor: 1 * Size:(1000, 100) | Memory: 0.3814 M | <class 'torch.Tensor'>]
memory_allocated:1.457520 Mb
max_memory_allocated:1.457520 Mb
memory_cached:2.000000 Mb
max_memory_cached:2.000000 Mb
Used Memory:9983.625000 Mb
Free Memory:1185.625000 Mb
Total Memory:11169.250000 Mb
0.0032593868672847748
LINE:41, FUNC:<module>, FILE:gpu_test.py, TIME:2019-06-02 17:16:23.791826, CONTENT:
[tensor: 1 * Size:(100, 10) | Memory: 0.0038 M | <class 'torch.Tensor'>]
[tensor: 1 * Size:() | Memory: 3.8146 M | <class 'torch.Tensor'>]
[tensor: 1 * Size:(64, 1000) | Memory: 0.2441 M | <class 'torch.Tensor'>]
[tensor: 2 * Size:(64, 10) | Memory: 0.0048 M | <class 'torch.Tensor'>]
[tensor: 1 * Size:(1000, 100) | Memory: 0.3814 M | <class 'torch.Tensor'>]
memory_allocated:1.846191 Mb
max_memory_allocated:2.252930 Mb
memory_cached:4.000000 Mb
max_memory_cached:4.000000 Mb
Used Memory:9989.625000 Mb
Free Memory:1179.625000 Mb
Total Memory:11169.250000 Mb

需要注意的是这里输出的大小根 nvidia-smi 展示的大小值不同。这里展示的值并没有包含pytorch gpu 上下文
占用的空间，所以即使，我代码里只有极少的gpu张量，nvidia-smi中显示的也有几百兆。

github 地址:

https://github.com/darr/pytorch_gpu_memory