Enable to register custom transform kernel
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 7.3k
- Avg merge
- 1d 15h
- Merged PRs (30d)
- 13
Description
🚀 The feature
It would be good to be able to register custom transform kernels in v2.function.
Motivation, pitch
If I want to register the transform's kernel, which is incompatible with built-in torchvision transforms and the functional API, and which uses built-in tv_tensor classes, it will be blocked by checking if it is from built-in function.
https://github.com/pytorch/vision/blob/6a9b5492d9590b19fe75300d95e3d9c4852a14ac/torchvision/transforms/v2/functional/_utils.py#L77-L84
https://github.com/pytorch/vision/blob/6a9b5492d9590b19fe75300d95e3d9c4852a14ac/torchvision/transforms/v2/functional/_utils.py#L92-L93
def custom_transform_kernel(inpt: torch.Tensor) -> torch.Tensor:
...
@F.register_kernel(custom_transform_kernel, tv_tensors.TVTensor)
def _custom_transform_kernel_dispatch(inpt: tv_tensors.TVTensor) -> tv_tensors.TVTensor:
output = custom_transform_kernel(inpt.as_subclass(torch.Tensor))
return tv_tensors.wrap(output, like=inpt)
class CustomTransform(tvt_v2.Transform):
def _transform(self, inpt: Any, params: dict[str, Any]) -> Any:
return self._call_kernel(custom_transform_kernel, inpt)
It would be more flexible if registering incompatible custom transform kernel is possible.
Alternatives
I tried to use @F._utils._register_kernel_internal instead, and it works.
But I think it could not be a safe way.
def custom_transform_kernel(inpt: torch.Tensor) -> torch.Tensor:
...
@F._utils._register_kernel_internal(custom_transform_kernel, tv_tensors.TVTensor, tv_tensor_wrapper=False)
def _custom_transform_kernel_dispatch(inpt: tv_tensors.TVTensor) -> tv_tensors.TVTensor:
output = custom_transform_kernel(inpt.as_subclass(torch.Tensor))
return tv_tensors.wrap(output, like=inpt)
class CustomTransform(tvt_v2.Transform):
def _transform(self, inpt: Any, params: dict[str, Any]) -> Any:
return self._call_kernel(custom_transform_kernel, inpt)
Additional context
No response
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with torchvision/transforms/v2/functional/_utils.py, especially the checks linked in the issue, and trace how v2.function and F.register_kernel handle custom kernels. Verify the example using a custom kernel and tv_tensors.TVTensor can register through the public API without the built-in-function restriction, while preserving the expected wrapped output.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- computer-vision
- Issue type
- Feature
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 38/100