基于mps的pytorch 多实例并行推理

# 服务开启
export CUDA_VISIBLE_DEVICES=0
nvidia-smi -i 0 -c 3 # -c 3 同 -c EXCLUSIVE_PROCESS
nvidia-cuda-mps-contrl -d

# 服务关闭
sudo nvidia-cuda-mps-control quit # 通常不能关闭mps服务，需要强制关闭
ps -ef | grep mps  # 查看mps的进程号
sudo kill -9 mps的进程号

import torch
import torchvision
import time
import os
import torchvision.models as models
from PIL import Image
import copy
device = torch.device('cuda')
# 装饰器，用于计算推理耗时
def getCompDuration(func):
    def warpper(*args, **kwargs):
        print("%s is running" % func.__name__)
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        computeTime = end - start
        print('forward time cost: %.5f sec' %computeTime)
        return result
    return wrapper
def getInputData(device = device):
    img = Image.open('your/data.jpg').convert('RGB')
    img = torchvision.transforms.ToTensor()(img)
    print(f'img:{img.shape}, rquires_grad:{img.requires_grad}')
    imgs = [img.to(device)]
    return imgs
@getCompDuration
def modelForwardImage(input, model, nTimes, stream, device = device):
    pid = os.getpid()
    model.eval()
    for i in range(nTimes):
        with torch.no_grad(): # 非常重要，用于降低显存消耗
            predictions = model(input)
            pred = predictions[0]['boxes'].shape
            print(f'pid:{pid}, stream:{stream}, predict result:{pred}')
def getModels(nProcess, device = device):
    modellist = []
  with torch.no_grad():
    model = models.detection.maskrcnn_resnet50_fpn(pretrained = False, pretrained_backbone = False)
        model.load_state_dict(torch.load('./your/dir/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth'))
        for i in range(nProcess):
            mod = copy.deepcopy(model)
            mod.to(device)
            modellist.append(mod)
        return modellist
def funcInStream(input, model, nTimes):
    stream = torch.cuda.Stream()
    with torch.cuda.stream(stream):
        modelForwardImage(input, model, nTimes, stream)
def test(nTimes, nProcess):
    input = getInputData()
    # spwan是多任务的一种方式
    ctx = torch.multiprocessing.get_context('spawn')
    models = getModels(nProcess)
    pool = []
    for i in range(nProcess):
        p = ctx.Process(target = funcInStream, args = (input, models[i], nTimes))
        pool.append(p)
    for p in pool:
        p.start()
if __name__ == '__main__':
    nTimes = 50
    nProcess = 2
    test(nTimes, nProcess)