Как надежно перехватить OOM графического процессора PyTorch в подпроцессе Python и динамически настроить пакетный размер ⇐ Python

[Anonymous]

Я создаю автономный агент искусственного интеллекта (управляющий рабочими процессами обучения), который автоматически генерирует сценарии обучения PyTorch/OpenMMLab и выполняет их в фоновом подпроцессе.
Одной из распространенных проблем, с которыми сталкивается мой агент, является RuntimeError: CUDA не хватает памяти. В настоящее время агент считывает stderr из subprocess.Popen, анализирует трассировку стека и перезаписывает файл конфигурации, чтобы уменьшить размер пакета вдвое.
Однако подпроцесс иногда оставляет зомби-процессы CUDA или невысвобожденную память на графическом процессоре, что приводит к немедленному сбою последующей автоматической повторной попытки.
Текущий упрощенный подход:

Код: Выделить всё

import subprocess
import re

def run_training_agent.py(config_path):
process = subprocess.Popen(["python", "train.py", config_path], stderr=subprocess.PIPE)
_, stderr = process.communicate()

if "CUDA out of memory" in stderr.decode('utf-8'):
# Agent logic to rewrite batch_size in config_path
reduce_batch_size(config_path)
# Restart
run_training_agent(config_path)

Мои вопросы:

• Каков наиболее надежный способ гарантировать, что исходный подпроцесс полностью освобождает VRAM до того, как агент инициирует повторную попытку? Должен ли я использовать многопроцессорность Python вместо подпроцесса?
• Есть ли способ перехватить исключение OOM внутри скрипта PyTorch с помощью перехватчика и сообщить точный объем памяти, необходимый обратно родительскому процессу агента, а не просто анализировать необработанный текст stderr?

Любые идеи по созданию устойчивых автоматизированных конвейеров машинного обучения будут полезны оценил!

Подробнее: https://stackoverflow.com/questions/799 ... ynamically