PyTorch tensor.to(device) for a List of Dict

tommyc: <p>I am working on an image object detection application using PyTorch <a href="https://pytorch.org/docs/stable/torchvision/models.html#torchvision.models.detection.fasterrcnn_resnet50_fpn" rel="nofollow noopener">torchvision.models.detection.fasterrcnn_resnet50_fpn</a>. As indicated by the documentation, during training phase, the input to fasterrcnn_resnet50_fpn model should be:</p>
<pre><code class="lang-auto">- 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
</code></pre>
<p>I am expecting training in mini-batches, so there should be more than 1 image tensor and 1 target dict in each list.</p>
<p>With the two lists prepared, I can train with my fasterrcnn_resnet50_fpn model:</p>
<pre><code class="lang-auto">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 ...

</code></pre>
<p>Everything works well, but there is one issue: by default, the tensor stored in cpu but I would like to train with gpu.</p>
<p>Naively, I can apply iterate through the lists and apply .to(device) for each tensor:</p>
<pre><code class="lang-auto">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 ]
</code></pre>
<p>However, I doubt it would be the most memory/computation efficient, as I am only looping through the lists.</p>
<p>Therefore, my question is: are there better methods to apply tensor.to(device) to list of tensors, or list of dicts, with better memory/computation efficiency, and probably better readability?</p>

albanD: <p>Hi,</p>
<p>I’m afraid there is None <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=9" title=":confused:" class="emoji" alt=":confused:"><br>
You can use something like this to make it cleaner but that’s pretty much it:</p>
<pre><code class="lang-python">def move_to(obj, device):
  if torch.is_tensor(obj):
    return obj.to(device):
  elif isinstance(obj, dict):
    res = {}
    for k, v in obj.items():
      res[k] = move_to(v, device)
    return res
  elif isinstance(obj, list):
    res = []
    for v in obj:
      res.append(move_to(v, device))
    return res
  else:
    raise TypeError("Invalid type for move_to")
</code></pre>
<pre><code class="lang-auto"></code></pre>

tommyc: <p>It would be nice if there is such a function with PyTorch. Your solution looks good. Thank you so much!</p>

Fei_Liu: <p>You can do this as a “quicker” workaround:</p>
<pre><code class="lang-auto">buffer = io.BytesIO()
torch.save(obj, buffer)
buffer.seek(0)
obj_cpu = torch.load(buffer, map_location='cpu')
</code></pre>
<p>This will make the tensors in <code>obj</code> all to cpu (with the new variable name <code>obj_cpu</code>). It works pretty much for whatever pickle-able by pickle/torch.</p>