Skip to content

Remove cuda dependency from Deformable-DETR #20993

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
alaradirik wants to merge 6 commits into from
Closed

Remove cuda dependency from Deformable-DETR #20993

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
fbfe93a
remove cuda dependency
alaradirik Jan 4, 2023
008acf0
Merge branch 'huggingface:main' into remove-cuda-deformabledetr
alaradirik Jan 4, 2023
a67c86c
remove redundant msda function
alaradirik Jan 4, 2023
638fd7b
revert changes, make custom cuda kernel optional
alaradirik Jan 5, 2023
881dc63
Merge branch 'huggingface:main' into remove-cuda-deformabledetr
alaradirik Jan 5, 2023
0ce0df0
fix merge conflicts
alaradirik Jan 5, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 0 additions & 51 deletions src/transformers/models/deformable_detr/load_custom.py
View file
Open in desktop

This file was deleted.

83 changes: 3 additions & 80 deletions src/transformers/models/deformable_detr/modeling_deformable_detr.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,6 @@
import torch
import torch.nn.functional as F
from torch import Tensor, nn
from torch.autograd import Function
from torch.autograd.function import once_differentiable

from ...activations import ACT2FN
from ...file_utils import (
Expand All @@ -34,84 +32,20 @@
add_start_docstrings_to_model_forward,
is_scipy_available,
is_timm_available,
is_torch_cuda_available,
is_vision_available,
replace_return_docstrings,
requires_backends,
)
from ...modeling_outputs import BaseModelOutput
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import meshgrid
from ...utils import is_ninja_available, logging
from ...utils import logging
from .configuration_deformable_detr import DeformableDetrConfig
from .load_custom import load_cuda_kernels


logger = logging.get_logger(__name__)

# Move this to not compile only when importing, this needs to happen later, like in __init__.
if is_torch_cuda_available() and is_ninja_available():
logger.info("Loading custom CUDA kernels...")
try:
MultiScaleDeformableAttention = load_cuda_kernels()
except Exception as e:
logger.warning(f"Could not load the custom kernel for multi-scale deformable attention: {e}")
MultiScaleDeformableAttention = None
else:
MultiScaleDeformableAttention = None

if is_vision_available():
from transformers.image_transforms import center_to_corners_format


class MultiScaleDeformableAttentionFunction(Function):
@staticmethod
def forward(
context,
value,
value_spatial_shapes,
value_level_start_index,
sampling_locations,
attention_weights,
im2col_step,
):
context.im2col_step = im2col_step
output = MultiScaleDeformableAttention.ms_deform_attn_forward(
value,
value_spatial_shapes,
value_level_start_index,
sampling_locations,
attention_weights,
context.im2col_step,
)
context.save_for_backward(
value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights
)
return output

@staticmethod
@once_differentiable
def backward(context, grad_output):
(
value,
value_spatial_shapes,
value_level_start_index,
sampling_locations,
attention_weights,
) = context.saved_tensors
grad_value, grad_sampling_loc, grad_attn_weight = MultiScaleDeformableAttention.ms_deform_attn_backward(
value,
value_spatial_shapes,
value_level_start_index,
sampling_locations,
attention_weights,
grad_output,
context.im2col_step,
)

return grad_value, None, None, grad_sampling_loc, grad_attn_weight, None


if is_scipy_available():
from scipy.optimize import linear_sum_assignment

Expand Down Expand Up @@ -664,19 +598,8 @@ def forward(
)
else:
raise ValueError(f"Last dim of reference_points must be 2 or 4, but got {reference_points.shape[-1]}")
try:
# GPU
output = MultiScaleDeformableAttentionFunction.apply(
value,
spatial_shapes,
level_start_index,
sampling_locations,
attention_weights,
self.im2col_step,
)
except Exception:
# CPU
output = ms_deform_attn_core_pytorch(value, spatial_shapes, sampling_locations, attention_weights)

output = ms_deform_attn_core_pytorch(value, spatial_shapes, sampling_locations, attention_weights)

@NielsRogge NielsRogge Jan 4, 2023

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So we're just switching to the CPU implementation here?

The Deformable DETR authors recommended to only use this for debugging:

The cuda version is much faster and more memory efficient than this pytorch version.

See thread: fundamentalvision/Deformable-DETR#9

@sgugger sgugger Jan 4, 2023

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh I initially thought it was a new function that was vectorized, not the CPU implementation. Let's not remove it indeed.

@shivalikasingh95 shivalikasingh95 Jan 4, 2023 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In that case, can't I use the CUDA version for Mask2Former as well? In fact even OneFormer could benefit from it if it's okay to use the CUDA implementation of MultiScaleDeformableAttention.

@NielsRogge NielsRogge Jan 4, 2023

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cc'ing @praeclarumjj3 (OneFormer author) here.

It would probably be best to provide the custom CUDA kernels as an option, rather than by default. This way, we can run the model in Google Colab, which is not the case at the moment.

We did the same for YOSO. This model has a boolean attribute in the config to determine whether or not to use the custom kernel.

@shivalikasingh95 shivalikasingh95 Jan 5, 2023

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@NielsRogge Agreed, it makes sense to keep it as optional. I can take up this change today for Mask2Former if you and @alaradirik think it makes sense to go ahead and add support for the CUDA implementation too. Let me know.

@alaradirik alaradirik Jan 5, 2023 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is not the cpu implementation, it can still be run on GPU. The PR simply removes the custom CUDA kernels, which parallelizes some of the GPU operations but also prevent users with an incompatible cuda library version to run the model on GPU.

@alaradirik alaradirik Jan 5, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

With that said, adding a boolean attribute to use the custom kernel is a good idea but users would need to first download the configuration of the checkpoint, edit it and then initialize the model if it's not going to be the default behaviour.

@NielsRogge @sgugger

@NielsRogge NielsRogge Jan 5, 2023

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm normally if you just do:

from transformers import DeformableDetrForObjectDetection

model = DeformableDetrForObjectDetection.from_pretrained("SenseTime/deformable-detr", use_custom_kernel=False)

then it should work. By providing additional kwargs to the from_pretrained method, you can edit the configuration.

@NielsRogge NielsRogge Jan 5, 2023 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

And @shivalikasingh95 sure but let's add support for the custom kernel in a separate PR for Mask2Former and OneFormer. We can set it to False by default, and add a boolean attribute for those models.

For Deformable DETR, we could change the behaviour to set it to False by default and inform users that if they really want to use the kernel, set the boolean attribute to True.

@sgugger sgugger Jan 5, 2023

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@alaradirik The function will be a lot slower on GPU since it's not parallelized, so enabling the custom CUDA kernel when possible would be better.

I don't think a flag is needed in the config. Let's just not fail if trying to load the CUDA kernel doesn't work and default to this non-vectorized implementation.

output = self.output_proj(output)

return output, attention_weights
Expand Down