PyTorch выделяет больше памяти на первом доступном графическом процессоре (cuda: 0) - PullRequest
Новая
       41

PyTorch выделяет больше памяти на первом доступном графическом процессоре (cuda: 0)

3 голосов
/ 23 января 2020

Как часть системы обучения с усиленным обучением, я параллельно тренирую четыре политики с использованием четырех графических процессоров. Для каждой модели есть два процесса - актер и ученик, которые используют только свой специфицированный c GPU (например, актер и ученик, соответствующие модели # 2, используют GPU # 2 только для всех своих тензоров). Актер и ученик делятся слоями модели через факел share_memory_(). Поскольку четыре обучающие «подсистемы» полностью симметричны c, я ожидаю, что они будут использовать одинаковый объем памяти графического процессора на каждом из четырех графических процессоров. На практике, однако, я вижу намного больше памяти GPU, выделенной на первом GPU (cuda:0).

Такое ощущение, что все совместное использование памяти каким-то образом выполняется через GPU # 0. Есть ли способ исправить это?

До сих пор я пытался установить CUDA_VISIBLE_DEVICES в дочерних процессах, явно изменив os.environ в функции процесса start. Похоже, что это не имеет никакого эффекта, возможно, потому, что дочерние процессы разветвляются из основного процесса, где PyTorch CUDA уже инициализирован, а envvars просто игнорируется в этот момент.

enter image description here

1 Ответ

3 голосов
/ 23 января 2020

Хорошо, пока я нашел обходной путь. Моя гипотеза была верна, если подсистема PyTorch CUDA уже инициализирована до того, как дочерний процесс разветвлен, установка CUDA_VISIBLE_DEVICES в другое значение для подпроцесса ничего не делает.

Хуже того, для инициализации достаточно вызова torch.cuda.device_count() CUDA, поэтому мы не можем даже запросить количество графических процессоров из PyTorch. Решением является либо жесткое его кодирование, передача в качестве параметра, либо запрос PyTorch API в отдельном процессе. Моя реализация для последнего: 

import sys


def get_available_gpus_without_triggering_pytorch_cuda_initialization(envvars):
    import subprocess
    out = subprocess.run([sys.executable, '-m', 'utils.get_available_gpus'], capture_output=True, env=envvars)
    text_output = out.stdout.decode()
    from utils.utils import log
    log.debug('Queried available GPUs: %s', text_output)
    return text_output


def main():
    import torch
    device_count = torch.cuda.device_count()
    available_gpus = ','.join(str(g) for g in range(device_count))
    print(available_gpus)
    return 0


if __name__ == '__main__':
    sys.exit(main())

По сути, эта функция вызывает свой собственный скрипт как отдельный python процесс и читает стандартный вывод.

Я не буду отмечать этот ответ как принятый, потому что я бы хотел бы узнать правильное решение, если оно существует.

...