Aluode/PerceptionLabPortable

1/*!2**************************************************************************************************3* Deformable DETR4* Copyright (c) 2020 SenseTime. All Rights Reserved.5* Licensed under the Apache License, Version 2.0 [see LICENSE for details]6**************************************************************************************************7* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.08**************************************************************************************************9*/10 11#pragma once12 13#include "cpu/ms_deform_attn_cpu.h"14 15#ifdef WITH_CUDA16#include "cuda/ms_deform_attn_cuda.h"17#endif18 19 20at::Tensor21ms_deform_attn_forward(22    const at::Tensor &value, 23    const at::Tensor &spatial_shapes,24    const at::Tensor &level_start_index,25    const at::Tensor &sampling_loc,26    const at::Tensor &attn_weight,27    const int im2col_step)28{29    if (value.type().is_cuda())30    {31#ifdef WITH_CUDA32        return ms_deform_attn_cuda_forward(33            value, spatial_shapes, level_start_index, sampling_loc, attn_weight, im2col_step);34#else35        AT_ERROR("Not compiled with GPU support");36#endif37    }38    AT_ERROR("Not implemented on the CPU");39}40 41std::vector<at::Tensor>42ms_deform_attn_backward(43    const at::Tensor &value, 44    const at::Tensor &spatial_shapes,45    const at::Tensor &level_start_index,46    const at::Tensor &sampling_loc,47    const at::Tensor &attn_weight,48    const at::Tensor &grad_output,49    const int im2col_step)50{51    if (value.type().is_cuda())52    {53#ifdef WITH_CUDA54        return ms_deform_attn_cuda_backward(55            value, spatial_shapes, level_start_index, sampling_loc, attn_weight, grad_output, im2col_step);56#else57        AT_ERROR("Not compiled with GPU support");58#endif59    }60    AT_ERROR("Not implemented on the CPU");61}62