Question

我发现在一个python进程中加载更多模型比在不同的python进程中加载每个模型消耗的GPU内存少得多。

假设以下代码在两个python进程中执行：

import torch
import time
import torchvision.models as models

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = models.densenet161(pretrained=False).to(device)

print('loaded')

time.sleep(1000)

nvidia-smi输出为：

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0     12546      C   python                                       773MiB |
|    0     13218      C   python                                       773MiB |
+-----------------------------------------------------------------------------+

这是合理的，因为我们已经加载了两个模型。但是，如果我在一个进程中加载两个模型，即运行以下代码：

import torch
import time
import torchvision.models as models

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = models.densenet161(pretrained=False).to(device)
model2 = models.densenet161(pretrained=False).to(device)

print('loaded')

time.sleep(1000)

nvidia-smi输出显示它消耗的内存比以前的情况少得多。

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0     13448      C   python                                       885MiB |
+-----------------------------------------------------------------------------+

pytorch是否使用了某种缓存重用技术？附言我也在张量流模型中观察到了这种现象。

问题是，我有很多要部署的pytorch模型，每个模型都需要大量的GPU内存。我不想将所有这些模型放在一个服务器进程中。我想使用docker容器管理每个模型。但是，如果我将模型部署在不同的Docker容器中，那么这些模型将消耗比我所能提供的更多的GPU内存。我的情况如何解决？

GPU内存与pytorch和tensorflow中的模型数量不成比例地增加

0 个答案: