Unwrap features before passing them into a kernel by pmeier · Pull Request #6807 · pytorch/vision

pmeier · 2022-10-20T21:57:41Z

Closes #6781 by going with the self.as_subclass(torch.Tensor) proposal. After this PR, there are no more __torch_function__ calls in the whole pipeline.

In fact, it almost renders #6681 obsolete: ndim, device, and dtype are no longer accessed. Even shape is only accessed 4 times:

vision/torchvision/prototype/transforms/functional/_meta.py

Lines 19 to 22 in 246de07

    
           elif isinstance(image, (features.Image, features.Video)): 
        
               channels = image.num_channels 
        
               height, width = image.spatial_size 
        
               return [channels, height, width]

vision/torchvision/prototype/transforms/functional/_meta.py

Lines 40 to 41 in 246de07

    
           elif isinstance(image, (features.Image, features.Video)): 
        
               return image.num_channels

vision/torchvision/prototype/transforms/functional/_meta.py

Lines 78 to 79 in 246de07

    
           elif isinstance(inpt, (features.Image, features.Video, features.BoundingBox, features.Mask)): 
        
               return list(inpt.spatial_size)

vision/torchvision/prototype/transforms/functional/_meta.py

Lines 91 to 92 in 246de07

    
           elif isinstance(inpt, features.Video): 
        
               return inpt.num_frames

Since they all only indirectly access .shape through properties we have anyway

vision/torchvision/prototype/features/_image.py

Lines 105 to 111 in 246de07

    
           @property 
        
           def spatial_size(self) -> Tuple[int, int]: 
        
               return cast(Tuple[int, int], tuple(self.shape[-2:])) 
        
           @property 
        
           def num_channels(self) -> int: 
        
               return self.shape[-3]

vision/torchvision/prototype/features/_video.py

Lines 57 to 67 in 246de07

    
           @property 
        
           def spatial_size(self) -> Tuple[int, int]: 
        
               return cast(Tuple[int, int], tuple(self.shape[-2:])) 
        
           @property 
        
           def num_channels(self) -> int: 
        
               return self.shape[-3] 
        
           @property 
        
           def num_frames(self) -> int: 
        
               return self.shape[-4]

we could also opt to move the with DisableTorchFunction() context manager there and completely remove

vision/torchvision/prototype/features/_feature.py

Lines 132 to 152 in 246de07

    
           # Add properties for common attributes like shape, dtype, device, ndim etc 
        
           # this way we return the result without passing into __torch_function__ 
        
           @property 
        
           def shape(self) -> _size:  # type: ignore[override] 
        
               with DisableTorchFunction(): 
        
                   return super().shape 
        
           @property 
        
           def ndim(self) -> int:  # type: ignore[override] 
        
               with DisableTorchFunction(): 
        
                   return super().ndim 
        
           @property 
        
           def device(self, *args: Any, **kwargs: Any) -> _device:  # type: ignore[override] 
        
               with DisableTorchFunction(): 
        
                   return super().device 
        
           @property 
        
           def dtype(self) -> _dtype:  # type: ignore[override] 
        
               with DisableTorchFunction(): 
        
                   return super().dtype

cc @vfdev-5 @datumbox @bjuncek

torchvision/prototype/features/_image.py

torchvision/prototype/features/_bounding_box.py

vfdev-5 · 2022-10-21T08:08:00Z

@pmeier can you provide time measurements on a datapipeline (e.g. image classification + AA + RandomErasing, #6681 (comment))

pmeier · 2022-10-21T08:44:04Z

$ python -u main.py cprofile_pil_vs_feature --n=2000

main

         1538052 function calls (1523449 primitive calls) in 2.556 seconds

   Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
      309    0.276    0.001    0.276    0.001 {built-in method torch.grid_sampler}
     2000    0.268    0.000    0.268    0.000 {built-in method torch._C._nn._upsample_bilinear2d_aa}
     9909    0.211    0.000    0.263    0.000 {method 'to' of 'torch._C._TensorBase' objects}
    16747    0.206    0.000    0.255    0.000 /home/philip/git/pytorch/torchvision/torchvision/prototype/features/_feature.py:63(__torch_function__)
     2348    0.091    0.000    0.091    0.000 {built-in method torch.round}
1499/1199    0.082    0.000    0.083    0.000 {built-in method torch.where}
     4039    0.055    0.000    0.055    0.000 {method 'clone' of 'torch._C._TensorBase' objects}
     2000    0.048    0.000    1.008    0.001 /home/philip/git/pytorch/torchvision/torchvision/prototype/transforms/_auto_augment.py:284(forward)
     2000    0.044    0.000    0.069    0.000 /home/philip/git/pytorch/torchvision/torchvision/transforms/functional_tensor.py:938(erase)
      899    0.042    0.000    0.042    0.000 {method 'scatter_add_' of 'torch._C._TensorBase' objects}

PR

         1380923 function calls (1366923 primitive calls) in 2.432 seconds

 Ordered by: internal time

 ncalls  tottime  percall  cumtime  percall filename:lineno(function)
    309    0.279    0.001    0.279    0.001 {built-in method torch.grid_sampler}
   2000    0.273    0.000    0.273    0.000 {built-in method torch._C._nn._upsample_bilinear2d_aa}
   9909    0.238    0.000    0.238    0.000 {method 'to' of 'torch._C._TensorBase' objects}
   2348    0.090    0.000    0.090    0.000 {built-in method torch.round}
   1199    0.084    0.000    0.084    0.000 {built-in method torch.where}
   4039    0.063    0.000    0.063    0.000 {method 'clone' of 'torch._C._TensorBase' objects}
   5798    0.047    0.000    0.047    0.000 {method 'div_' of 'torch._C._TensorBase' objects}
   2000    0.045    0.000    0.957    0.000 /home/philip/git/pytorch/torchvision/torchvision/prototype/transforms/_auto_augment.py:284(forward)
   2000    0.045    0.000    0.045    0.000 {method 'flip' of 'torch._C._TensorBase' objects}
   2000    0.043    0.000    0.072    0.000 /home/philip/git/pytorch/torchvision/torchvision/transforms/functional_tensor.py:938(erase)

The diff that we are interested in:

-     16747    0.206    0.000    0.255    0.000 /home/philip/git/pytorch/torchvision/torchvision/prototype/features/_feature.py:63(__torch_function__)
-     14487    0.022    0.000    0.022    0.000 {method 'as_subclass' of 'torch._C._TensorBase' objects}
+     18896    0.032    0.000    0.032    0.000 {method 'as_subclass' of 'torch._C._TensorBase' objects}

Meaning we are looking at a (0.206 + 0.022 - 0.032) / 2000 roughly 100µs improvement per sample.

vfdev-5 · 2022-10-21T08:47:39Z

@pmeier by time measurements I meant torch benchmarks and not cprofile (which may be less reliable)

datumbox

LGTM overall. Only one comment below. Also mypy seems to be failing due to an unused ignore.

torchvision/prototype/features/_feature.py

datumbox

LGTM, thanks!

Summary: * unwrap features before calling the kernels * revert double unwrapping * cleanup * remove debug raise * more cleanup Reviewed By: YosuaMichael Differential Revision: D40588165 fbshipit-source-id: 3ce277bbe4f47124f572ca8cd185795b5917fe3e

pmeier added 2 commits October 20, 2022 17:02

unwrap features before calling the kernels

10e87ef

Merge branch 'main' into unwrap

b89a3bd

pmeier added module: transforms Perf For performance improvements prototype labels Oct 20, 2022

pmeier requested review from datumbox and vfdev-5 October 20, 2022 21:57

facebook-github-bot added the cla signed label Oct 20, 2022

pmeier commented Oct 20, 2022

View reviewed changes

torchvision/prototype/features/_image.py Show resolved Hide resolved

datumbox reviewed Oct 21, 2022

View reviewed changes

torchvision/prototype/features/_bounding_box.py Outdated Show resolved Hide resolved

revert double unwrapping

48fc5ba

cleanup

f13d7eb

datumbox reviewed Oct 21, 2022

View reviewed changes

torchvision/prototype/features/_feature.py Outdated Show resolved Hide resolved

pmeier added 3 commits October 21, 2022 11:28

remove debug raise

977a837

more cleanup

071079a

Merge branch 'main' into unwrap

c8102ef

datumbox approved these changes Oct 21, 2022

View reviewed changes

Merge branch 'main' into unwrap

367a840

pmeier merged commit 3761855 into pytorch:main Oct 21, 2022

pmeier deleted the unwrap branch October 21, 2022 11:34

pmeier mentioned this pull request Oct 25, 2022

unwrap features in dispatchers #6831

Merged

pmeier mentioned this pull request Aug 7, 2023

Preserve Datapoint subclasses instead of returning tensors #7807

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Unwrap features before passing them into a kernel#6807

Unwrap features before passing them into a kernel#6807
pmeier merged 8 commits intopytorch:mainfrom
pmeier:unwrap

pmeier commented Oct 20, 2022 •

edited by pytorch-bot bot

Loading

Uh oh!

Uh oh!

Uh oh!

vfdev-5 commented Oct 21, 2022

Uh oh!

pmeier commented Oct 21, 2022

Uh oh!

vfdev-5 commented Oct 21, 2022

Uh oh!

datumbox left a comment

Uh oh!

Uh oh!

datumbox left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

	elif isinstance(image, (features.Image, features.Video)):
	channels = image.num_channels
	height, width = image.spatial_size
	return [channels, height, width]

	elif isinstance(image, (features.Image, features.Video)):
	return image.num_channels

	elif isinstance(inpt, (features.Image, features.Video, features.BoundingBox, features.Mask)):
	return list(inpt.spatial_size)

	elif isinstance(inpt, features.Video):
	return inpt.num_frames

	@property
	def spatial_size(self) -> Tuple[int, int]:
	return cast(Tuple[int, int], tuple(self.shape[-2:]))

	@property
	def num_channels(self) -> int:
	return self.shape[-3]

	# Add properties for common attributes like shape, dtype, device, ndim etc
	# this way we return the result without passing into __torch_function__
	@property
	def shape(self) -> _size: # type: ignore[override]
	with DisableTorchFunction():
	return super().shape

	@property
	def ndim(self) -> int: # type: ignore[override]
	with DisableTorchFunction():
	return super().ndim

	@property
	def device(self, args: Any, *kwargs: Any) -> _device: # type: ignore[override]
	with DisableTorchFunction():
	return super().device

	@property
	def dtype(self) -> _dtype: # type: ignore[override]
	with DisableTorchFunction():
	return super().dtype

Conversation

pmeier commented Oct 20, 2022 • edited by pytorch-bot bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

vfdev-5 commented Oct 21, 2022

Uh oh!

pmeier commented Oct 21, 2022

Uh oh!

vfdev-5 commented Oct 21, 2022

Uh oh!

datumbox left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

datumbox left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

pmeier commented Oct 20, 2022 •

edited by pytorch-bot bot

Loading