关于pytorch指定了gpu还是占用了别的gpu的显存的问题

问题如下：

**
今天搭了一个服务，用到了多个模型，我的4块显卡显存都是12G的，而所有的模型运行起来一共占用9G，所以制定了所有的模型运行环境都是 cuda:0 ，然而在运行的时候，总会有一块别的GPU被莫名占用。找了半天找到问题，记录一哈
**

1、onnx服务选择gpu

我的服务中用到了多个模型，有的是转成了onnx的，有的就是pth的，转换失败了。
一开始我以为是某个onnx模型没有指定，所以莫名加载到了别的gpu上，所以在我所有的onnx运行时中都加上了gpu指定。

代码如下：

import onnxruntime as ort
self.encoder = ort.InferenceSession(encoder_path,providers=['CUDAExecutionProvider'])
self.encoder.set_providers(['CUDAExecutionProvider'], [ {'device_id': 0}])
self.encoder_input_name = self.encoder.get_inputs()[0].name
features = self.encoder.run(None,{self.encoder_input_name:img})[0]

这里 device_id : 0 指的是第一块gpu

测试结果确实可以指定gpu，但是那个莫名占用的gpu显存的情况还是存在，看来是torch的模型的问题

2、torch 选择gpu

class Provincer():
    def __init__(self,module_path=r'checkpoint/epoch_100.pth',device='cuda:0'):
        self.device = torch.device(device)

        self.net = resnet18()
        self.net.load_state_dict(torch.load(module_path))
        self.net.to(device)
        self.net.eval()

一顿debug，终于，在province模型调用的时候发现占用。具体的是在torch.load 时第三块gpu显存开始占用。

原因不知道，个人认为，应该时训练的时候放在了第三块gpu上，所以默认加载会放到那里，然后再复制到我指定的gpu上。
所以这里修改一下默认加载位置为cpu。

class Provincer():
    def __init__(self,module_path=r'checkpoint/epoch_100.pth',device='cuda:0'):
        self.device = torch.device(device)

        self.net = resnet18()
        self.net.load_state_dict(torch.load(module_path,map_location='cpu'))
        self.net.to(device)
        self.net.eval()

重新加载，ok，果然可以了。

3、运行时设置

当然，这里还有个比较好的解决方法，运行时设置使用的gpu设备

如开始运行命令为

python3 server.py

可以修改为

CUDA_VISIBLE_DEVICES=1 python3 server.py

或者在代码中增加

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '1'

然后在搜索时发现有人提供了一种方法

将所有的加载代码放在with中，此方法没有尝试，先记录下来。

with torch.cuda.device(1):
	pass