GridRaySampler should not persist the grid.

[PeterL1n]

🚀 Feature

GridRaySampler creates the grid in constructor and calls register_buffer. However, we should pass persistent=False argument so that the canonical grid is not saved in state_dict.

pytorch3d/pytorch3d/renderer/implicit/raysampling.py

Line 99 in 8041178

self.register_buffer("_xy_grid", _xy_grid)

Motivation

We want the grid to be generated in code every time instead of loading from the state_dict. This allows us to change the image render size at the construction time. Currently, the state_dict will mismatch if the image size is changed. This is a needed feature since it is common to have training and evaluation to be at different resolutions.

Pitch

Change

self.register_buffer("_xy_grid", _xy_grid)

to

self.register_buffer("_xy_grid", _xy_grid, False)

BC Breaking!

Note: This will break existing state_dicts as the new module no longer has _xy_grid, but previously saved state_dicts will have the extra key. A simple workaround is to use load_state_dict(state_dict, strict=False) What are your thoughts on BC breaking?

[bottler]

I agree that this would be a good enhancement. We're thinking about the BC-breaking implications internally.

[PeterL1n]

@bottler
Some approaches:

1. Add another arg in the constructor persist_grid=True by default to be BC compatible. Then people who need interoperability can set it to False to prevent the grid from being persisted.
2. Override load_state_dict to safely ignore the grid if the newly generated is the same as before. Throw if the old grid is different. Also can print a message to tell the user that we detected a persisted grid from old versions. Please save the state_dict one more time to remove persisted grid. Future versions will no longer be compatible with the persisted grid.

[bottler]

Asking people to set strict=False is no more bothersome than asking them to change the code which creates their model. I'm wondering if there is a way to have a custom deserializer/unpickler in our code for the GridRaySampler which explicitly ignores any provided _xy_grid, which I think would be the nicest.

[PeterL1n]

@bottler Maybe something like this.

def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
      del state_dict[prefix + '_xy_grid']
      super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)

[bottler]

Yes. That is one of several ways to fix the problem, relying on internal parts of PyTorch. Another is to register a hook with self._register_load_state_dict_pre_hook.

We have decided that we prefer to avoid such things and just have a documented BC break in PyTorch3D. I plan to make the change soon.

[bottler]

This change is now committted.