Модель PyTorch занимает слишком много времени для загрузки в первый раз на новой машине

#python #amazon-ec2 #pytorch

Вопрос:

У меня есть настройка ручного масштабирования на EC2, где я создаю экземпляры на основе 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.