如何在python中持久化一个进程

Question

我有以下 Pytorch 代码，代码的作用是：

1. 使用net=CRAFT() 将 Pytorch 对象检测模型加载到 RAM 中
2. 使用net.eval() 运行推理
3. 我使用 for 循环提供图像以逐一运行推理

代码：

def copyStateDict(state_dict):
    if list(state_dict.keys())[0].startswith("module"):
        start_idx = 1
    else:
        start_idx = 0
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = ".".join(k.split(".")[start_idx:])
        new_state_dict[name] = v
    return new_state_dict


if __name__ == '__main__':
    # load net
    net = CRAFT()     # initialize

    print('Loading weights from checkpoint (' + args.trained_model + ')')
    if args.cuda:
        net.load_state_dict(copyStateDict(torch.load(args.trained_model)))
    else:
        net.load_state_dict(copyStateDict(torch.load(args.trained_model, map_location='cpu')))

    if args.cuda:
        net = net.cuda()
        net = torch.nn.DataParallel(net)
        cudnn.benchmark = False

    net.eval()

    t = time.time()

    # load data
    for k, image_path in enumerate(image_list):
        print("Test image {:d}/{:d}: {:s}".format(k+1, len(image_list), image_path), end='\r')
        image = imgproc.loadImage(image_path)

        bboxes, polys, score_text = test_net(net, image, args.text_threshold, args.link_threshold, args.low_text, args.cuda, args.poly)

        # save score text
        #filename, file_ext = os.path.splitext(os.path.basename(image_path))
        #mask_file = result_folder + "/res_" + filename + '_mask.jpg'
        #cv2.imwrite(mask_file, score_text)

        file_utils.saveResult(image_path, image[:,:,::-1], polys, dirname=result_folder)

    print("elapsed time : {}s".format(time.time() - t))

这段代码有一个警告，尤其是我每次执行步骤 1 和 2 时它都会占用1.8 GB 的 RAM。对图像的实际推断需要更少的 RAM（每张图像 40mb），但我的大部分内存消耗来自将 ML 模型加载到 RAM 中。有没有一种方法可以让我将此模型继续加载到服务器上的 RAM 中，然后仅对图像执行推断？

Answer 1

您几乎描述了一个 Web 服务器。您可以使用 flask 之类的东西围绕此创建一个简单的 API，它将在启动时加载模型并在 Web 请求中接收图像或路径，然后返回响应。使用uwsgi 或类似的方式运行它，一次有多个工作人员可用（只需确保模型初始化一次，而不是每个工作人员）。

或者您可以更花哨并使用众多 RPC 库之一来创建类似的本地服务器。