PyTorch模型在新机器上第一次加载时花费太多



我在EC2上有一个手动缩放设置,我正在创建基于AMI的实例,该AMI已经在启动时运行我的代码(使用Systemd)。我面临着一个基本问题:在主实例(我用来创建AMI的实例)上,Python代码在映像启动后需要8秒才能准备好,这包括导入库、加载模型的状态字典等。现在,在我用AMI创建的映像上,代码第一次启动需要5分钟以上的时间,将状态字典从磁盘加载到GPU内存需要特别长的时间,在第一次之后,代码需要与加载主实例相同的时间。

AMI保持相同的pycache文件夹作为主实例,所以应该不会花太多时间,因为我认为AMI应该包含所有内容,不是吗?所以,我的问题是:是否有任何其他缓存使CUDA/Python更快,我没有考虑到?我只保留pycache/文件夹,但我不知道是否有什么我能做的,以确保它不会花费那么多时间来启动所有的第一次。这是我的主要结构:
# Import libraries
import torch
import numpy as np
# Import personal models (takes 1 minute)
from model1 import model1
from model2 import model2
# Load first model
model1_object = model1()
model2_object = model2()
# Load state dicts (takes 3+ minutes, the first time in new instances, seconds other times)
# Note: the models are a bit heavy
model1_object.load_state_dict(torch.load("model1.pth"))
model2_object.load_state_dict(torch.load("model2.pth"))

注意:我使用g4dn。xlarge实例,用于主实例和AWS中的新实例。

这是由于恢复AWS EBS快照所需的高延迟造成的。首先,当您恢复快照时,延迟非常高,这解释了为什么在我的示例中,当实例是新创建的时候,加载模型需要花费这么多时间。

查看本文的初始化部分:https://cloudonaut.io/ebs-snapshot-pitfalls/

我找到的唯一解决方案是在实例第一次创建时快速使用它是启用快速快照恢复,这花费大约500美元一个月:https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ebs-fast-snapshot-restore.html

如果您有空闲时间,您可以等到达到最佳性能,或者尝试事先预热音量https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ebs-initialize.html

最新更新