Тупик с torch.distributed.rp c с num_workers> 1 - PullRequest
Новая
       9

Тупик с torch.distributed.rp c с num_workers> 1

0 голосов
/ 19 февраля 2020

У меня есть большой (93 ГБ) файл .h5, содержащий функции изображений в моей локальной системе, и моя модель обучается на кластере SLURM ADA, который имеет ограничение хранения 25 ГБ. Я пытаюсь использовать torch.distributed.rp c framework для запроса функций изображения в Dataset.getitem с использованием удаленного вызова сервера rp c в моей локальной системе.

Код для инициализации сервера RP C (локальная система): 

import os
import torch.distributed.rpc as rpc

def run_worker(rank, world_size):
    os.environ['MASTER_ADDR'] = 'XX.X.XX.XX'
    os.environ['MASTER_PORT'] = 'XXXX'

    rpc.init_rpc(utils.SERVER_NAME,
                 rank=rank, 
                 world_size=world_size)
    print("Server Initialized", flush=True)

    rpc.shutdown()


if __name__ == "__main__":
    world_size = 2
    rank = 0

    run_worker(rank, world_size)

Код для сервера RP C для запроса данных из локальной системы (в ADA), 

import os
import torch.distributed.rpc as rpc

def run_worker(rank, world_size):
    os.environ['MASTER_ADDR'] = 'XX.X.XX.XX'
    os.environ['MASTER_PORT'] = 'XXXX'

    rpc.init_rpc(utils.CLIENT_NAME.format(rank), rank=rank, world_size=world_size)
    print("Client Initialized", flush=True)

    main()

    rpc.shutdown()

if __name__ == '__main__':
    world_size = 2
    rank = 1

    run_worker(rank, world_size)

В моем data_loader я указал num_worker = 8 ,
Упрощенный код для dataset.getitem: (на ADA), 

def __getitem__(self, index):
        ....
        ....

        print("fetching image for image_id {}, item {}".format(image_id, item), flush=True)
        v = utils._remote_method(utils.Server._load_image, self.server_ref, [index, self.coco_id_to_index])

        return v, ......

Теперь в моем обучении l oop, когда я вызываю enumerate (data_loader) , многопроцессная загрузка данных включена и getitem функция вызывается num_worker раз и тупик достигнут. Я не уверен, почему возникает эта тупиковая ситуация, потому что всякий раз, когда вызывается getitem , необходимо выполнить удаленный вызов на сервер RP C в моей локальной системе для запроса данных.

Как я могу разрешить тупик? Есть ли другой способ решить мою проблему, если большой файл не помещается в мою систему ADA, я не хочу идти на компромисс с моей задержкой.

Редактировать: тупик разрешается, когда я устанавливаю num_workers = 0 , но он очень медленный, 20 секунд / итерации.

...