Pytorch3d: Camera as Attribute of nn.Module not pushed to GPU by .to(device)

[florianHofherr]

🐛 Bugs / Unexpected behaviors

Hi everyone

I encounter the following unexpected behavior. I have a nn.Module class which has a pytorch3d camera as an attribute.

import torch
import torch.nn as nn
from pytorch3d.renderer import PerspectiveCameras

class test_cam_to_GPU(nn.Module):
  def __init__(self):
    super().__init__()

    self.linear = nn.Linear(3, 4)
    self.register_buffer('buffer', torch.rand(3))
    self.cam = PerspectiveCameras()

I don't want to optimize over the camera parameters, I just want it to be part of the model. When I now move an instance of test_cam_to_GPU to the GPU, then the .cam attribute stays on the CPU. I would expect it to be treated like the buffer, which is indeed moved to the GPU. I am not quite sure if this behavior is intended, to me it is somewhat unexpected, in particular since the camera object inherits from nn.Module, and even has a custom .to(device) method.

I have already looked into it a little bit, since I assumed that the .to(device) method is called recursively on all the child objects. But that does not seem to be the case. As far as I have understood, the method is only called once, and for every child object, all attributes of a certain type (parameters, buffers,...) are moved to the device. Since the camera parameters do not belong to any of the classes, they are not moved.

If this behavior should be changed, I think it is sufficient to do minor modifications to TensorProperties. Potentially it is enough to change the setattr to register_buffer (if the attribute is not a parameter, probably there needs to be a check for that).

I appreciate any other ideas, maybe I have missed something and this is not an issue at all.

Quick fix

In case anyone has the same issue, I solved it by implementing a custom .to(device) method, that manually moves the camera.
Seems to work, but I would prefer this to work without the manual adjustments

class custom_move_method(nn.Module):
    def __init__(self):
        super().__init__()

        self.linear = nn.Linear(3, 4)
        self.register_buffer('buffer', torch.rand(3))
        self.cam = PerspectiveCameras()

    def to(self, device = 'cpu'):
        self = super(custom_move_method, self).to(device)
        self.cam.to(device)
        return self

[gkioxari]

Hi @florianHofherr
Thank you for bringing to our attention and for the quick fix suggestion. I labelled it as "potential-bug" so that we incorporate it into the upcoming library fixes.

[nikhilaravi]

@florianHofherr Thanks for flagging this. We did look into this briefly previously but will revisit this issue. I think the problem arises because within TensorProperties we may modify the input (i.e. convert to tensor if the input is a list/tuple and/or broadcast some dimensions.

You are correct that the calling .to on a module does not recursively call .toon all the child modules so we need a different solution.

[florianHofherr]

Thanks for looking into it