PyTorch tenor.to (устройство) для списка Dict - PullRequest
Новая
       94

PyTorch tenor.to (устройство) для списка Dict

1 голос
/ 10 января 2020

Я работаю над приложением обнаружения объектов изображения, используя PyTorch torchvision.models.detection.fasterrcnn_resnet50_fpn . Как указано в документации, на этапе обучения вход для модели fastrcnn_resnet50_fpn должен быть: 

- list of image tensors, each of shape [C, H, W]
- list of target dicts, each with:
    - boxes (FloatTensor[N, 4]): the ground-truth boxes in [x1, y1, x2, y2] format, 
                                 with values between 0 and H and 0 and W
    - labels (Int64Tensor[N]): the class label for each ground-truth box

Я ожидаю обучения в мини-пакетах, поэтому должно быть более 1 тензора изображения и 1 целевой диктат в каждый список.

С двумя подготовленными списками я могу тренироваться с моей более быстрой моделью rcnn_resnet50_fpn: 

model.train()
for image_list, target_list in dataset_loader:
    """
    image_list: list of image tensors
    target_list: list of dicts {boxes, labels}
    """
    # some steps before ...

    # Feed inputs to model in training phase
    outputs = model(image_list, target_list)

    # more steps after ...

Все работает хорошо, но есть одна проблема: по умолчанию тензор хранится в процессоре но я бы хотел тренироваться с gpu.

Наивно, я могу применить итерацию по спискам и применить .to (устройство) для каждого тензора: 

image_list = [ t.to(device) for t in image_list ]
target_list = [ {'boxes':d['boxes'].to(device), 'labels':d['labels']} for d in target_list ]

Однако я сомневаюсь, что это будет быть максимально эффективным для памяти / вычислений, так как я только перебираю списки.

Поэтому мой вопрос таков: существуют ли более эффективные методы для применения tenor.to (устройства) к списку тензоров или списку слов с большей эффективностью использования памяти / вычислений и, возможно, лучшей читабельностью?

...