<p dir="auto">I want to perform on the GPU non-maximum-suppression on the output of a darknet/yolo CNN.<br>
The following testing code works fine :</p>
<p dir="auto">//Credits: adapted from <a href="https://github.com/pprp">https://github.com/pprp</a> to test nms parallel algo with at::Tensor as input<br>
#include &lt;stdio.h&gt;<br>
#include &lt;stdlib.h&gt;<br>
#include "string"<br>
//#include &lt;windows.h&gt;//needed for LoadLibrary<br>
#include &lt;cuda_runtime.h&gt;<br>
#include <br>
#include &lt;opencv2/core/core.hpp&gt;<br>
#include &lt;opencv2/highgui/highgui.hpp&gt;<br>
#include "opencv2/imgproc/imgproc.hpp"<br>
#include "device_launch_parameters.h"</p>
<p dir="auto">#include "device_functions.h"</p>
<p dir="auto">#include &lt;ATen/ATen.h&gt;<br>
#include &lt;ATen/cuda/CUDAContext.h&gt;<br>
#include &lt;THC/THC.h&gt;<br>
#include &lt;THC/THCDeviceUtils.cuh&gt;</p>
<p dir="auto">#include "torch/torch.h"//-&gt;including torch gives problems!!!???</p>
<p dir="auto">using namespace std;</p>
<p dir="auto">#define HANDLE_ERROR(ans) { gpuAssert((ans), <strong>FILE</strong>, <strong>LINE</strong>); }</p>
<p dir="auto">inline void gpuAssert(cudaError_t code, const char *file, int line, bool abort=true)<br>
{<br>
if (code != cudaSuccess)<br>
{<br>
fprintf(stderr,"GPUassert: %s %s %d\n", cudaGetErrorString(code), file, line);<br>
if (abort) exit(code);<br>
}<br>
}</p>
<p dir="auto">typedef struct<br>
{<br>
double x,y,w,h;<br>
char s[100];<br>
char cls[100];<br>
double cmps;<br>
}box;</p>
<p dir="auto"><strong>device</strong> inline float devIoU(float const* const   b1, float const* const  b2) {<br>
float ai = (float)(b1[2] + 1) * (b1[3] + 1);<br>
float aj = (float)(b2[2] + 1) * (b2[3] + 1);<br>
float x_inter, x2_inter, y_inter, y2_inter;</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="x_inter = max(b1[0], b2[0]);
y_inter = max(b1[1], b2[1]);

x2_inter = min((b1[0] + b1[2]), (b2[0] + b2[2]));
y2_inter = min((b1[1] + b1[3]), (b2[1] + b2[3]));

float w = (float)max((float)0, x2_inter - x_inter);
float h = (float)max((float)0, y2_inter - y_inter);

float inter = ((w * h) / (ai + aj - w * h));
return inter;"><pre class="notranslate"><code class="notranslate">x_inter = max(b1[0], b2[0]);
y_inter = max(b1[1], b2[1]);

x2_inter = min((b1[0] + b1[2]), (b2[0] + b2[2]));
y2_inter = min((b1[1] + b1[3]), (b2[1] + b2[3]));

float w = (float)max((float)0, x2_inter - x_inter);
float h = (float)max((float)0, y2_inter - y_inter);

float inter = ((w * h) / (ai + aj - w * h));
return inter;
</code></pre></div>
<p dir="auto">}</p>
<p dir="auto"><strong>global</strong> void NMS_GPU(const int n_boxes, const float nms_overlap_thresh,<br>
const float* dev_boxes, bool* d_res) {<br>
unsigned int xIndex = blockIdx.x * blockDim.x + threadIdx.x;<br>
//unsigned int xIndex = threadIdx.x;//only 1 block with index 0!<br>
float cur_box[5];<br>
float a_box[5];<br>
cur_box[0] = dev_boxes[xIndex * 5 + 0];<br>
cur_box[1] = dev_boxes[xIndex * 5 + 1];<br>
cur_box[2] = dev_boxes[xIndex * 5 + 2];<br>
cur_box[3] = dev_boxes[xIndex * 5 + 3];<br>
cur_box[4] = dev_boxes[xIndex * 5 + 4];<br>
//__syncthreads();//not necessary as cur_box is not a shared resource<br>
for (int i = 0; i &lt; 19; i++)<br>
{<br>
if (i != xIndex)<br>
{<br>
a_box[0] = dev_boxes[i * 5 + 0];<br>
a_box[1] = dev_boxes[i * 5 + 1];<br>
a_box[2] = dev_boxes[i * 5 + 2];<br>
a_box[3] = dev_boxes[i * 5 + 3];<br>
a_box[4] = dev_boxes[i * 5 + 4];<br>
if (a_box[4] &lt; cur_box[4] )<br>
{<br>
if (devIoU(a_box, cur_box) &gt; nms_overlap_thresh)<br>
{<br>
d_res[i] = false;<br>
}<br>
}<br>
}<br>
}<br>
}</p>
<p dir="auto">int main()<br>
{<br>
int const threadsPerBlock = sizeof(unsigned long long) * 8;//Bufo : =64(float size)<br>
//LoadLibrary(TEXT("D:\dev\Cpp\dependencies\torchnew\lib\torch_cuda.dll"));<br>
at::DeviceType device_type;</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="if (at::cuda::is_available()) {
	device_type = at::kCUDA;
}
else {
	device_type = at::kCPU;
	std::cout &lt;&lt; &quot;No GPU avalaible, sorry ...&quot; &lt;&lt; std::endl;
	return 0;
}
at::Device device(device_type);
std::cout &lt;&lt; &quot;Device : &quot; &lt;&lt; device_type &lt;&lt; std::endl;

const int count = 19;
cv::Mat temp = cv::imread(&quot;Cow_45.jpg&quot;,1);

bool *h_res =(bool*)malloc(sizeof(bool)*count);//contains the result of the nms algo (cpu context)
for(int i=0; i&lt;count; i++)
{
	h_res[i] = true;
}

box b[count];
b[0].x = 996.000000;b[0].y = 2566.420000;b[0].w = 170.793000;b[0].h=172.580000;
strcpy(b[0].cls,&quot;nose&quot;);strcpy(b[0].s,&quot;0.983194&quot;);b[0].cmps=0.983194;
b[1].x = 4238.937000;b[1].y = 1594.513000;b[1].w = 160.063000;b[1].h=148.487000;
strcpy(b[1].cls,&quot;eye&quot;);strcpy(b[1].s,&quot;0.992166&quot;);b[1].cmps=0.992166;
b[2].x = 4656.389000;b[2].y = 2175.186000;b[2].w = 316.180000;b[2].h=221.552000;
strcpy(b[2].cls,&quot;nose&quot;);strcpy(b[2].s,&quot;0.994816&quot;);b[2].cmps=0.994816;
b[3].x = 4316.000000;b[3].y = 1660.000000;b[3].w = 127.474000;b[3].h=113.452000;
strcpy(b[3].cls,&quot;eye&quot;);strcpy(b[3].s,&quot;0.990833&quot;);b[3].cmps=0.990833;
b[4].x = 997.013000;b[4].y = 2664.408000;b[4].w = 222.214000;b[4].h=229.068000;
strcpy(b[4].cls,&quot;nose&quot;);strcpy(b[4].s,&quot;0.985067&quot;);b[4].cmps=0.985067;
b[5].x = 666.069000;b[5].y = 2029.219000;b[5].w = 135.689000;b[5].h=160.833000;
strcpy(b[5].cls,&quot;eye&quot;);strcpy(b[5].s,&quot;0.993240&quot;);b[5].cmps=0.993240;
b[6].x = 4653.547000;b[6].y = 2324.000000;b[6].w = 338.125000;b[6].h=133.902000;
strcpy(b[6].cls,&quot;nose&quot;);strcpy(b[6].s,&quot;0.982858&quot;);b[6].cmps=0.982858;
b[7].x = 4476.556000;b[7].y = 2131.557000;b[7].w = 253.402000;b[7].h=273.601000;
strcpy(b[7].cls,&quot;nose&quot;);strcpy(b[7].s,&quot;0.959098&quot;);b[7].cmps=0.959098;
b[8].x = 754.326000;b[8].y = 2571.066000;b[8].w = 324.674000;b[8].h=161.605000;
strcpy(b[8].cls,&quot;nose&quot;);strcpy(b[8].s,&quot;0.993699&quot;);b[8].cmps=0.993699;
b[9].x = 729.962000;b[9].y = 2658.741000;b[9].w = 349.038000;b[9].h=192.046000;
strcpy(b[9].cls,&quot;nose&quot;);strcpy(b[9].s,&quot;0.986209&quot;);b[9].cmps=0.986209;
b[10].x = 1271.863000;b[10].y = 2058.679000;b[10].w = 138.781000;b[10].h=137.553000;
strcpy(b[10].cls,&quot;eye&quot;);strcpy(b[10].s,&quot;0.989965&quot;);b[10].cmps=0.989965;
b[11].x = 4316.000000;b[11].y = 1601.751000;b[11].w = 134.204000;b[11].h=141.249000;
strcpy(b[11].cls,&quot;eye&quot;);strcpy(b[11].s,&quot;0.988307&quot;);b[11].cmps=0.988307;
b[12].x = 650.901000;b[12].y = 2032.621000;b[12].w = 91.484000;b[12].h=42.112000;
strcpy(b[12].cls,&quot;eye&quot;);strcpy(b[12].s,&quot;0.969982&quot;);b[12].cmps=0.969982;
b[13].x = 1328.000000;b[13].y = 2058.692000;b[13].w = 103.849000;b[13].h=136.518000;
strcpy(b[13].cls,&quot;eye&quot;);strcpy(b[13].s,&quot;0.987316&quot;);b[13].cmps=0.987316;
b[14].x = 214.809000;b[14].y = 1599.809000;b[14].w = 1553.705000;b[14].h=1319.679000;
strcpy(b[14].cls,&quot;head&quot;);strcpy(b[14].s,&quot;0.997623&quot;);b[14].cmps=0.997623;
b[15].x = 3826.177000;b[15].y = 1072.206000;b[15].w = 1254.063000;b[15].h=1412.903000;
strcpy(b[15].cls,&quot;head&quot;);strcpy(b[15].s,&quot;0.997487&quot;);b[15].cmps=0.997487;
b[16].x = 729.632000;b[16].y = 2578.523000;b[16].w = 442.495000;b[16].h=302.378000;
strcpy(b[16].cls,&quot;nose&quot;);strcpy(b[16].s,&quot;0.960093&quot;);b[16].cmps=0.960093;
b[17].x = 655.430000;b[17].y = 2031.151000;b[17].w = 91.570000;b[17].h=148.691000;
strcpy(b[17].cls,&quot;eye&quot;);strcpy(b[17].s,&quot;0.993275&quot;);b[17].cmps=0.993275;
b[18].x = 4251.712000;b[18].y = 1660.000000;b[18].w = 147.288000;b[18].h=105.309000;
strcpy(b[18].cls,&quot;eye&quot;);strcpy(b[18].s,&quot;0.992576&quot;);b[18].cmps=0.992576;

//***************************************************************************************************************************************
//copy boxes to a  at::tensor
at::Tensor boxes = at::zeros({ count,5 });

for (int i = 0; i &lt; count; i++) {
	boxes[i][0] = b[i].x;
	boxes[i][1] = b[i].y;
	boxes[i][2] = b[i].w;
	boxes[i][3] = b[i].h;
	boxes[i][4] = b[i].cmps;
}

boxes  = boxes.to(device);
float* boxes_as_array = boxes.data&lt;float&gt;();//convert at::tensor to a flat array

float nms_overlap_thresh = 0.1;

//Comment: this piece of code is apparently necessary to assign the Torch cuda context to the global cuda context
THCState* state = at::globalContext().lazyInitCUDA(); // TODO replace with getTHCState
const int col_blocks = THCCeilDiv(count, threadsPerBlock);
unsigned long long* mask_dev = NULL;
mask_dev = (unsigned long long*) THCudaMalloc(state,count * col_blocks * sizeof(unsigned long long));
//-------------------------------------------------------------------------------------------------------------

bool *d_res;
//port h_res to GPU
HANDLE_ERROR(cudaMalloc((void**)&amp;d_res, count*sizeof(bool)));
HANDLE_ERROR(cudaMemcpy(d_res, h_res,sizeof(bool)*count, cudaMemcpyHostToDevice));

NMS_GPU&lt;&lt;&lt;dim3(1,count,1),count&gt;&gt;&gt;(count,nms_overlap_thresh, boxes_as_array,d_res);

//port d_res to CPU
HANDLE_ERROR(cudaMemcpy(h_res, d_res, sizeof(bool)*count, cudaMemcpyDeviceToHost));

//display result
for(int i =0; i&lt;count ; i++)
{
	if(*(h_res+i) == true)
	{
		//printf(&quot;GPU Draw: %d--%d\n&quot;,i,*(h_res+i));
		cv::putText(temp,b[i].cls,cv::Point((int)b[i].x,(int)b[i].y-5),cv::FONT_HERSHEY_SIMPLEX,1.7,cv::Scalar(255,255,255),5,8,0);
		cv::putText(temp,b[i].s,cv::Point((int)b[i].x+120,(int)b[i].y-5),cv::FONT_HERSHEY_SIMPLEX,1.7,cv::Scalar(255,255,255),5,8,0);
		cv::rectangle(temp,cv::Point((int)b[i].x,(int)b[i].y),cv::Point((int)b[i].x + (int)b[i].w,(int)b[i].y + (int)b[i].h),cv::Scalar(92.185,194),8,8,0);
	}
}
cv::namedWindow(&quot;Window&quot;,0);
cv::resizeWindow(&quot;Window&quot;,1064,800);
cv::imshow(&quot;Window&quot;,temp);
cv::waitKey(0);
return 0;"><pre class="notranslate"><code class="notranslate">if (at::cuda::is_available()) {
	device_type = at::kCUDA;
}
else {
	device_type = at::kCPU;
	std::cout &lt;&lt; "No GPU avalaible, sorry ..." &lt;&lt; std::endl;
	return 0;
}
at::Device device(device_type);
std::cout &lt;&lt; "Device : " &lt;&lt; device_type &lt;&lt; std::endl;

const int count = 19;
cv::Mat temp = cv::imread("Cow_45.jpg",1);

bool *h_res =(bool*)malloc(sizeof(bool)*count);//contains the result of the nms algo (cpu context)
for(int i=0; i&lt;count; i++)
{
	h_res[i] = true;
}

box b[count];
b[0].x = 996.000000;b[0].y = 2566.420000;b[0].w = 170.793000;b[0].h=172.580000;
strcpy(b[0].cls,"nose");strcpy(b[0].s,"0.983194");b[0].cmps=0.983194;
b[1].x = 4238.937000;b[1].y = 1594.513000;b[1].w = 160.063000;b[1].h=148.487000;
strcpy(b[1].cls,"eye");strcpy(b[1].s,"0.992166");b[1].cmps=0.992166;
b[2].x = 4656.389000;b[2].y = 2175.186000;b[2].w = 316.180000;b[2].h=221.552000;
strcpy(b[2].cls,"nose");strcpy(b[2].s,"0.994816");b[2].cmps=0.994816;
b[3].x = 4316.000000;b[3].y = 1660.000000;b[3].w = 127.474000;b[3].h=113.452000;
strcpy(b[3].cls,"eye");strcpy(b[3].s,"0.990833");b[3].cmps=0.990833;
b[4].x = 997.013000;b[4].y = 2664.408000;b[4].w = 222.214000;b[4].h=229.068000;
strcpy(b[4].cls,"nose");strcpy(b[4].s,"0.985067");b[4].cmps=0.985067;
b[5].x = 666.069000;b[5].y = 2029.219000;b[5].w = 135.689000;b[5].h=160.833000;
strcpy(b[5].cls,"eye");strcpy(b[5].s,"0.993240");b[5].cmps=0.993240;
b[6].x = 4653.547000;b[6].y = 2324.000000;b[6].w = 338.125000;b[6].h=133.902000;
strcpy(b[6].cls,"nose");strcpy(b[6].s,"0.982858");b[6].cmps=0.982858;
b[7].x = 4476.556000;b[7].y = 2131.557000;b[7].w = 253.402000;b[7].h=273.601000;
strcpy(b[7].cls,"nose");strcpy(b[7].s,"0.959098");b[7].cmps=0.959098;
b[8].x = 754.326000;b[8].y = 2571.066000;b[8].w = 324.674000;b[8].h=161.605000;
strcpy(b[8].cls,"nose");strcpy(b[8].s,"0.993699");b[8].cmps=0.993699;
b[9].x = 729.962000;b[9].y = 2658.741000;b[9].w = 349.038000;b[9].h=192.046000;
strcpy(b[9].cls,"nose");strcpy(b[9].s,"0.986209");b[9].cmps=0.986209;
b[10].x = 1271.863000;b[10].y = 2058.679000;b[10].w = 138.781000;b[10].h=137.553000;
strcpy(b[10].cls,"eye");strcpy(b[10].s,"0.989965");b[10].cmps=0.989965;
b[11].x = 4316.000000;b[11].y = 1601.751000;b[11].w = 134.204000;b[11].h=141.249000;
strcpy(b[11].cls,"eye");strcpy(b[11].s,"0.988307");b[11].cmps=0.988307;
b[12].x = 650.901000;b[12].y = 2032.621000;b[12].w = 91.484000;b[12].h=42.112000;
strcpy(b[12].cls,"eye");strcpy(b[12].s,"0.969982");b[12].cmps=0.969982;
b[13].x = 1328.000000;b[13].y = 2058.692000;b[13].w = 103.849000;b[13].h=136.518000;
strcpy(b[13].cls,"eye");strcpy(b[13].s,"0.987316");b[13].cmps=0.987316;
b[14].x = 214.809000;b[14].y = 1599.809000;b[14].w = 1553.705000;b[14].h=1319.679000;
strcpy(b[14].cls,"head");strcpy(b[14].s,"0.997623");b[14].cmps=0.997623;
b[15].x = 3826.177000;b[15].y = 1072.206000;b[15].w = 1254.063000;b[15].h=1412.903000;
strcpy(b[15].cls,"head");strcpy(b[15].s,"0.997487");b[15].cmps=0.997487;
b[16].x = 729.632000;b[16].y = 2578.523000;b[16].w = 442.495000;b[16].h=302.378000;
strcpy(b[16].cls,"nose");strcpy(b[16].s,"0.960093");b[16].cmps=0.960093;
b[17].x = 655.430000;b[17].y = 2031.151000;b[17].w = 91.570000;b[17].h=148.691000;
strcpy(b[17].cls,"eye");strcpy(b[17].s,"0.993275");b[17].cmps=0.993275;
b[18].x = 4251.712000;b[18].y = 1660.000000;b[18].w = 147.288000;b[18].h=105.309000;
strcpy(b[18].cls,"eye");strcpy(b[18].s,"0.992576");b[18].cmps=0.992576;

//***************************************************************************************************************************************
//copy boxes to a  at::tensor
at::Tensor boxes = at::zeros({ count,5 });

for (int i = 0; i &lt; count; i++) {
	boxes[i][0] = b[i].x;
	boxes[i][1] = b[i].y;
	boxes[i][2] = b[i].w;
	boxes[i][3] = b[i].h;
	boxes[i][4] = b[i].cmps;
}

boxes  = boxes.to(device);
float* boxes_as_array = boxes.data&lt;float&gt;();//convert at::tensor to a flat array

float nms_overlap_thresh = 0.1;

//Comment: this piece of code is apparently necessary to assign the Torch cuda context to the global cuda context
THCState* state = at::globalContext().lazyInitCUDA(); // TODO replace with getTHCState
const int col_blocks = THCCeilDiv(count, threadsPerBlock);
unsigned long long* mask_dev = NULL;
mask_dev = (unsigned long long*) THCudaMalloc(state,count * col_blocks * sizeof(unsigned long long));
//-------------------------------------------------------------------------------------------------------------

bool *d_res;
//port h_res to GPU
HANDLE_ERROR(cudaMalloc((void**)&amp;d_res, count*sizeof(bool)));
HANDLE_ERROR(cudaMemcpy(d_res, h_res,sizeof(bool)*count, cudaMemcpyHostToDevice));

NMS_GPU&lt;&lt;&lt;dim3(1,count,1),count&gt;&gt;&gt;(count,nms_overlap_thresh, boxes_as_array,d_res);

//port d_res to CPU
HANDLE_ERROR(cudaMemcpy(h_res, d_res, sizeof(bool)*count, cudaMemcpyDeviceToHost));

//display result
for(int i =0; i&lt;count ; i++)
{
	if(*(h_res+i) == true)
	{
		//printf("GPU Draw: %d--%d\n",i,*(h_res+i));
		cv::putText(temp,b[i].cls,cv::Point((int)b[i].x,(int)b[i].y-5),cv::FONT_HERSHEY_SIMPLEX,1.7,cv::Scalar(255,255,255),5,8,0);
		cv::putText(temp,b[i].s,cv::Point((int)b[i].x+120,(int)b[i].y-5),cv::FONT_HERSHEY_SIMPLEX,1.7,cv::Scalar(255,255,255),5,8,0);
		cv::rectangle(temp,cv::Point((int)b[i].x,(int)b[i].y),cv::Point((int)b[i].x + (int)b[i].w,(int)b[i].y + (int)b[i].h),cv::Scalar(92.185,194),8,8,0);
	}
}
cv::namedWindow("Window",0);
cv::resizeWindow("Window",1064,800);
cv::imshow("Window",temp);
cv::waitKey(0);
return 0;
</code></pre></div>
<p dir="auto">}</p>
<p dir="auto">Problem : When I include "torch/torch.h"  (uncomment line 23) - something I need for rest of the project -  I get the following error (Visual Studio 2019):</p>
<p dir="auto">"Error		member "torch::jit::detail::ParameterPolicy::all_slots" may not be initialized	acudaNMSTEST	D:\dev\Cpp\dependencies\torchnew\include\torch\csrc\jit\api\module.h	490	"</p>
<p dir="auto">I struggled a couple of days to find a solution, but in vain. Any idea what is going on/what I did wrong?</p>
<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/malfet/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/malfet">@malfet</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/yf225/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/yf225">@yf225</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/glaringlee/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/glaringlee">@glaringlee</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/peterjc123/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/peterjc123">@peterjc123</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/nbcsm/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/nbcsm">@nbcsm</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/guyang3532/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/guyang3532">@guyang3532</a></p>
<p dir="auto">win10 64 ,libtorch 1.5 release<br>
I tried to run the <code class="notranslate">psroiales.cu</code> file to report an error.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="member of &quot;the torch: : jit: : detail: : AttributePolicy: : all_slots&quot; may not be initialized
member &quot;torch::jit::detail::BufferPolicy::all_slots&quot; may not be initialized  
member &quot;torch::jit::detail::ModulePolicy::all_slots&quot; may not be initialized	
member &quot;torch::jit::detail::ParameterPolicy::all_slots&quot; may not be initialized	
D:\deeplearning\deeplearning\learn\libtorch\include\torch\csrc\jit\api\module.h	501	
错误	MSB3721	命令“&quot;C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.2\bin\nvcc.exe&quot;  --use-local-env -ccbin &quot;D:\vs2019\VC\Tools\MSVC\14.26.28801\bin\HostX64\x64&quot; -x cu  -I&quot;C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.2\include&quot; -I&quot;C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.2\include&quot;     --keep-dir x64\Release -maxrregcount=0  --machine 64 --compile -cudart static     -DNDEBUG -D_CONSOLE -D_UNICODE -DUNICODE -Xcompiler &quot;/EHsc /W3 /nologo /O2 /FdD:\deeplearning\x64\Release\dir\vc142.pdb /FS /Zi  /MD /GR&quot; -o D:\deeplearning\x64\Release\dir\PSROIAlign_cuda.cu.obj &quot;D:\deeplearning\deeplearning\PSROIAlign_cuda.cu&quot;”已退出，返回代码为 1。	deeplearning	D:\vs2019\MSBuild\Microsoft\VC\v160\BuildCustomizations\CUDA 10.2.targets	764	"><pre class="notranslate"><code class="notranslate">member of "the torch: : jit: : detail: : AttributePolicy: : all_slots" may not be initialized
member "torch::jit::detail::BufferPolicy::all_slots" may not be initialized  
member "torch::jit::detail::ModulePolicy::all_slots" may not be initialized	
member "torch::jit::detail::ParameterPolicy::all_slots" may not be initialized	
D:\deeplearning\deeplearning\learn\libtorch\include\torch\csrc\jit\api\module.h	501	
错误	MSB3721	命令“"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.2\bin\nvcc.exe"  --use-local-env -ccbin "D:\vs2019\VC\Tools\MSVC\14.26.28801\bin\HostX64\x64" -x cu  -I"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.2\include" -I"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.2\include"     --keep-dir x64\Release -maxrregcount=0  --machine 64 --compile -cudart static     -DNDEBUG -D_CONSOLE -D_UNICODE -DUNICODE -Xcompiler "/EHsc /W3 /nologo /O2 /FdD:\deeplearning\x64\Release\dir\vc142.pdb /FS /Zi  /MD /GR" -o D:\deeplearning\x64\Release\dir\PSROIAlign_cuda.cu.obj "D:\deeplearning\deeplearning\PSROIAlign_cuda.cu"”已退出，返回代码为 1。	deeplearning	D:\vs2019\MSBuild\Microsoft\VC\v160\BuildCustomizations\CUDA 10.2.targets	764	
</code></pre></div>
<details>
<summary>PSROIAlign_cuda.cu</summary>
<div class="highlight highlight-source-c notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="#include &lt;ATen/cuda/CUDAApplyUtils.cuh&gt;
#include &quot;PSROIAlign_cuda.h&quot;


template &lt;typename T&gt;
__device__ T bilinear_interpolate(
    const T* input,
    const int height,
    const int width,
    T y,
    T x,
    const int index /* index for debug only*/) {
  // deal with cases that inverse elements are out of feature map boundary
  if (y &lt; -1.0 || y &gt; height || x &lt; -1.0 || x &gt; width) {
    // empty
    return 0;
  }

  if (y &lt;= 0)
    y = 0;
  if (x &lt;= 0)
    x = 0;

  int y_low = (int)y;
  int x_low = (int)x;
  int y_high;
  int x_high;

  if (y_low &gt;= height - 1) {
    y_high = y_low = height - 1;
    y = (T)y_low;
  } else {
    y_high = y_low + 1;
  }

  if (x_low &gt;= width - 1) {
    x_high = x_low = width - 1;
    x = (T)x_low;
  } else {
    x_high = x_low + 1;
  }

  T ly = y - y_low;
  T lx = x - x_low;
  T hy = 1. - ly, hx = 1. - lx;

  // do bilinear interpolation
  T v1 = input[y_low * width + x_low];
  T v2 = input[y_low * width + x_high];
  T v3 = input[y_high * width + x_low];
  T v4 = input[y_high * width + x_high];
  T w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;

  T val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);

  return val;
}

template &lt;typename T&gt;
__global__ void PSROIAlignForwardCUDA(
    const int nthreads,
    const T* input,
    const T spatial_scale,
    const int channels,
    const int height,
    const int width,
    const int pooled_height,
    const int pooled_width,
    const int sampling_ratio,
    const T* rois,
    const int channels_out,
    T* output,
    int* channel_mapping) {
  CUDA_1D_KERNEL_LOOP(index, nthreads) {
    // (n, c_out, ph, pw) is an element in the pooled output
    int pw = index % pooled_width;
    int ph = (index / pooled_width) % pooled_height;
    int c_out = (index / pooled_width / pooled_height) % channels_out;
    int n = index / pooled_width / pooled_height / channels_out;

    // (n, c_in, ph, pw) is the associated element in the input
    int c_in = (c_out * pooled_height + ph) * pooled_width + pw;

    // [start, end) interval for spatial sampling
    const T* offset_rois = rois + n * 5;
    int roi_batch_ind = offset_rois[0];

    // Do not using rounding; this implementation detail is critical
    T roi_start_w = offset_rois[1] * spatial_scale - static_cast&lt;T&gt;(0.5);
    T roi_start_h = offset_rois[2] * spatial_scale - static_cast&lt;T&gt;(0.5);
    T roi_end_w = offset_rois[3] * spatial_scale - static_cast&lt;T&gt;(0.5);
    T roi_end_h = offset_rois[4] * spatial_scale - static_cast&lt;T&gt;(0.5);

    T roi_width = roi_end_w - roi_start_w;
    T roi_height = roi_end_h - roi_start_h;
    T bin_size_h = static_cast&lt;T&gt;(roi_height) / static_cast&lt;T&gt;(pooled_height);
    T bin_size_w = static_cast&lt;T&gt;(roi_width) / static_cast&lt;T&gt;(pooled_width);

    // Do not using floor/ceil; this implementation detail is critical
    T hstart = static_cast&lt;T&gt;(ph) * bin_size_h + roi_start_h;
    T wstart = static_cast&lt;T&gt;(pw) * bin_size_w + roi_start_w;

    // We use roi_bin_grid to sample the grid and mimic integral
    int roi_bin_grid_h = (sampling_ratio &gt; 0)
        ? sampling_ratio
        : ceil(roi_height / pooled_height);
    int roi_bin_grid_w =
        (sampling_ratio &gt; 0) ? sampling_ratio : ceil(roi_width / pooled_width);
    const T count = roi_bin_grid_h * roi_bin_grid_w;

    const T* offset_input =
        input + (roi_batch_ind * channels + c_in) * height * width;
    T out_sum = 0;
    for (int iy = 0; iy &lt; roi_bin_grid_h; iy++) {
      const T y = hstart +
          static_cast&lt;T&gt;(iy + .5f) * bin_size_h /
              static_cast&lt;T&gt;(roi_bin_grid_h);
      for (int ix = 0; ix &lt; roi_bin_grid_w; ix++) {
        const T x = wstart +
            static_cast&lt;T&gt;(ix + .5f) * bin_size_w /
                static_cast&lt;T&gt;(roi_bin_grid_w);
        T val = bilinear_interpolate(offset_input, height, width, y, x, index);
        out_sum += val;
      }
    }

    out_sum /= count;
    output[index] = out_sum;
    channel_mapping[index] = c_in;
  }
}

template &lt;typename T&gt;
__device__ void bilinear_interpolate_gradient(
    const int height,
    const int width,
    T y,
    T x,
    T&amp; w1,
    T&amp; w2,
    T&amp; w3,
    T&amp; w4,
    int&amp; x_low,
    int&amp; x_high,
    int&amp; y_low,
    int&amp; y_high,
    const int index /* index for debug only*/) {
  // deal with cases that inverse elements are out of feature map boundary
  if (y &lt; -1.0 || y &gt; height || x &lt; -1.0 || x &gt; width) {
    // empty
    w1 = w2 = w3 = w4 = 0.;
    x_low = x_high = y_low = y_high = -1;
    return;
  }

  if (y &lt;= 0)
    y = 0;
  if (x &lt;= 0)
    x = 0;

  y_low = (int)y;
  x_low = (int)x;

  if (y_low &gt;= height - 1) {
    y_high = y_low = height - 1;
    y = (T)y_low;
  } else {
    y_high = y_low + 1;
  }

  if (x_low &gt;= width - 1) {
    x_high = x_low = width - 1;
    x = (T)x_low;
  } else {
    x_high = x_low + 1;
  }

  T ly = y - y_low;
  T lx = x - x_low;
  T hy = 1. - ly, hx = 1. - lx;

  // reference in forward
  // T v1 = input[y_low * width + x_low];
  // T v2 = input[y_low * width + x_high];
  // T v3 = input[y_high * width + x_low];
  // T v4 = input[y_high * width + x_high];
  // T val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);

  w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;

  return;
}

template &lt;typename T&gt;
__global__ void PSROIAlignBackwardCUDA(
    const int nthreads,
    const T* grad_output,
    const int* channel_mapping,
    const int num_rois,
    const T spatial_scale,
    const int channels,
    const int height,
    const int width,
    const int pooled_height,
    const int pooled_width,
    const int sampling_ratio,
    const int channels_out,
    T* grad_input,
    const T* rois) {
  CUDA_1D_KERNEL_LOOP(index, nthreads) {
    // (n, *, ph, pw) is an element in the pooled output
    int pw = index % pooled_width;
    int ph = (index / pooled_width) % pooled_height;
    int n = index / pooled_width / pooled_height / channels_out;

    const T* offset_rois = rois + n * 5;
    int roi_batch_ind = offset_rois[0];

    // Do not using rounding; this implementation detail is critical
    T roi_start_w = offset_rois[1] * spatial_scale - static_cast&lt;T&gt;(0.5);
    T roi_start_h = offset_rois[2] * spatial_scale - static_cast&lt;T&gt;(0.5);
    T roi_end_w = offset_rois[3] * spatial_scale - static_cast&lt;T&gt;(0.5);
    T roi_end_h = offset_rois[4] * spatial_scale - static_cast&lt;T&gt;(0.5);

    // Force too small ROIs to be 1x1
    T roi_width = roi_end_w - roi_start_w;
    T roi_height = roi_end_h - roi_start_h;
    T bin_size_h = roi_height / static_cast&lt;T&gt;(pooled_height);
    T bin_size_w = roi_width / static_cast&lt;T&gt;(pooled_width);

    int c_in = channel_mapping[index];
    T* grad_input_offset =
        grad_input + (roi_batch_ind * channels + c_in) * height * width;

    // Do not using floor/ceil; this implementation detail is critical
    T hstart = static_cast&lt;T&gt;(ph) * bin_size_h + roi_start_h;
    T wstart = static_cast&lt;T&gt;(pw) * bin_size_w + roi_start_w;

    const T grad_output_this_bin = grad_output[index];

    // We use roi_bin_grid to sample the grid and mimic integral
    int roi_bin_grid_h = (sampling_ratio &gt; 0)
        ? sampling_ratio
        : ceil(roi_height / pooled_height); // e.g., = 2
    int roi_bin_grid_w =
        (sampling_ratio &gt; 0) ? sampling_ratio : ceil(roi_width / pooled_width);
    const T count = roi_bin_grid_h * roi_bin_grid_w;

    for (int iy = 0; iy &lt; roi_bin_grid_h; iy++) {
      const T y = hstart +
          static_cast&lt;T&gt;(iy + .5f) * bin_size_h /
              static_cast&lt;T&gt;(roi_bin_grid_h);
      for (int ix = 0; ix &lt; roi_bin_grid_w; ix++) {
        const T x = wstart +
            static_cast&lt;T&gt;(ix + .5f) * bin_size_w /
                static_cast&lt;T&gt;(roi_bin_grid_w);

        T w1, w2, w3, w4;
        int x_low, x_high, y_low, y_high;

        bilinear_interpolate_gradient(
            height,
            width,
            y,
            x,
            w1,
            w2,
            w3,
            w4,
            x_low,
            x_high,
            y_low,
            y_high,
            index);

        T g1 = grad_output_this_bin * w1 / count;
        T g2 = grad_output_this_bin * w2 / count;
        T g3 = grad_output_this_bin * w3 / count;
        T g4 = grad_output_this_bin * w4 / count;

        if (x_low &gt;= 0 &amp;&amp; x_high &gt;= 0 &amp;&amp; y_low &gt;= 0 &amp;&amp; y_high &gt;= 0) {
          atomicAdd(grad_input_offset + y_low * width + x_low, g1);
          atomicAdd(grad_input_offset + y_low * width + x_high, g2);
          atomicAdd(grad_input_offset + y_high * width + x_low, g3);
          atomicAdd(grad_input_offset + y_high * width + x_high, g4);
        } // if
      } // ix
    } // iy
  }
}

std::tuple&lt;at::Tensor, at::Tensor&gt; PSROIAlign_forward_cuda(
    const at::Tensor&amp; input,
    const at::Tensor&amp; rois,
    const float spatial_scale,
    const int pooled_height,
    const int pooled_width,
    const int sampling_ratio) {
  // Check if input tensors are CUDA tensors
  AT_ASSERTM(input.type().is_cuda(), &quot;input must be a CUDA tensor&quot;);
  AT_ASSERTM(rois.type().is_cuda(), &quot;rois must be a CUDA tensor&quot;);

  at::TensorArg input_t{input, &quot;input&quot;, 1}, rois_t{rois, &quot;rois&quot;, 2};

  at::CheckedFrom c = &quot;PSROIAlign_forward_cuda&quot;;
  at::checkAllSameGPU(c, {input_t, rois_t});
  at::checkAllSameType(c, {input_t, rois_t});

  at::cuda::CUDAGuard device_guard(input.device());

  auto num_rois = rois.size(0);
  auto channels = input.size(1);
  auto height = input.size(2);
  auto width = input.size(3);

  AT_ASSERTM(
      channels % (pooled_height * pooled_width) == 0,
      &quot;input channels must be a multiple of pooling height * pooling width&quot;);
  int channels_out = channels / (pooled_height * pooled_width);

  auto output = at::zeros(
      {num_rois, channels_out, pooled_height, pooled_width}, input.options());
  auto channel_mapping =
      at::zeros(output.sizes(), input.options().dtype(at::kInt));

  auto output_size = output.numel();
  if (output_size == 0) {
    AT_CUDA_CHECK(cudaGetLastError());
    return std::make_tuple(output, channel_mapping);
  }

  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  dim3 grid(std::min(
      at::cuda::ATenCeilDiv(
          static_cast&lt;int64_t&gt;(output_size), static_cast&lt;int64_t&gt;(512)),
      static_cast&lt;int64_t&gt;(4096)));
  dim3 block(512);

  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
      input.scalar_type(), &quot;PSROIAlign_forward&quot;, [&amp;] {
        PSROIAlignForwardCUDA&lt;scalar_t&gt;&lt;&lt;&lt;grid, block, 0, stream&gt;&gt;&gt;(
            output_size,
            input.contiguous().data_ptr&lt;scalar_t&gt;(),
            spatial_scale,
            channels,
            height,
            width,
            pooled_height,
            pooled_width,
            sampling_ratio,
            rois.contiguous().data_ptr&lt;scalar_t&gt;(),
            channels_out,
            output.data_ptr&lt;scalar_t&gt;(),
            channel_mapping.data_ptr&lt;int&gt;());
      });
  AT_CUDA_CHECK(cudaGetLastError());
  cudaDeviceSynchronize();
  return std::make_tuple(output, channel_mapping);
}

at::Tensor PSROIAlign_backward_cuda(
    const at::Tensor&amp; grad,
    const at::Tensor&amp; rois,
    const at::Tensor&amp; channel_mapping,
    const float spatial_scale,
    const int pooled_height,
    const int pooled_width,
    const int sampling_ratio,
    const int batch_size,
    const int channels,
    const int height,
    const int width) {
  // Check if input tensors are CUDA tensors
  AT_ASSERTM(grad.type().is_cuda(), &quot;grad must be a CUDA tensor&quot;);
  AT_ASSERTM(rois.type().is_cuda(), &quot;rois must be a CUDA tensor&quot;);
  AT_ASSERTM(
      channel_mapping.type().is_cuda(),
      &quot;channel_mapping must be a CUDA tensor&quot;);

  at::TensorArg grad_t{grad, &quot;grad&quot;, 1}, rois_t{rois, &quot;rois&quot;, 2},
      channel_mapping_t{channel_mapping, &quot;channel_mapping&quot;, 3};

  at::CheckedFrom c = &quot;PSROIAlign_backward_cuda&quot;;
  at::checkAllSameGPU(c, {grad_t, rois_t, channel_mapping_t});
  at::checkAllSameType(c, {grad_t, rois_t});

  at::cuda::CUDAGuard device_guard(grad.device());

  auto num_rois = rois.size(0);
  auto grad_input =
      at::zeros({batch_size, channels, height, width}, grad.options());

  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  dim3 grid(std::min(
      at::cuda::ATenCeilDiv(
          static_cast&lt;int64_t&gt;(grad.numel()), static_cast&lt;int64_t&gt;(512)),
      static_cast&lt;int64_t&gt;(4096)));
  dim3 block(512);

  // handle possibly empty gradients
  if (grad.numel() == 0) {
    AT_CUDA_CHECK(cudaGetLastError());
    return grad_input;
  }

  int channels_out = channels / (pooled_height * pooled_width);

  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
      grad.scalar_type(), &quot;PSROIAlign_backward&quot;, [&amp;] {
        PSROIAlignBackwardCUDA&lt;scalar_t&gt;&lt;&lt;&lt;grid, block, 0, stream&gt;&gt;&gt;(
            grad.numel(),
            grad.contiguous().data_ptr&lt;scalar_t&gt;(),
            channel_mapping.data_ptr&lt;int&gt;(),
            num_rois,
            spatial_scale,
            channels,
            height,
            width,
            pooled_height,
            pooled_width,
            sampling_ratio,
            channels_out,
            grad_input.data_ptr&lt;scalar_t&gt;(),
            rois.contiguous().data_ptr&lt;scalar_t&gt;());
      });
  AT_CUDA_CHECK(cudaGetLastError());
  return grad_input;
}"><pre class="notranslate"><span class="pl-k">#include</span> <span class="pl-s">&lt;ATen/cuda/CUDAApplyUtils.cuh&gt;</span>
<span class="pl-k">#include</span> <span class="pl-s">"PSROIAlign_cuda.h"</span>


<span class="pl-s1">template</span> <span class="pl-c1">&lt;</span><span class="pl-s1">typename</span> <span class="pl-c1">T</span><span class="pl-c1">&gt;</span>
<span class="pl-s1">__device__</span> <span class="pl-smi">T</span> <span class="pl-en">bilinear_interpolate</span>(
    <span class="pl-k">const</span> <span class="pl-smi">T</span><span class="pl-c1">*</span> <span class="pl-s1">input</span>,
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">height</span>,
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">width</span>,
    <span class="pl-smi">T</span> <span class="pl-s1">y</span>,
    <span class="pl-smi">T</span> <span class="pl-s1">x</span>,
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">index</span> <span class="pl-c">/* index for debug only*/</span>) {
  <span class="pl-c">// deal with cases that inverse elements are out of feature map boundary</span>
  <span class="pl-k">if</span> (<span class="pl-s1">y</span> <span class="pl-c1">&lt;</span> <span class="pl-c1">-1.0</span> <span class="pl-c1">||</span> <span class="pl-s1">y</span> <span class="pl-c1">&gt;</span> <span class="pl-s1">height</span> <span class="pl-c1">||</span> <span class="pl-s1">x</span> <span class="pl-c1">&lt;</span> <span class="pl-c1">-1.0</span> <span class="pl-c1">||</span> <span class="pl-s1">x</span> <span class="pl-c1">&gt;</span> <span class="pl-s1">width</span>) {
    <span class="pl-c">// empty</span>
    <span class="pl-k">return</span> <span class="pl-c1">0</span>;
  }

  <span class="pl-k">if</span> (<span class="pl-s1">y</span> &lt;= <span class="pl-c1">0</span>)
    <span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>;
  <span class="pl-k">if</span> (<span class="pl-s1">x</span> &lt;= <span class="pl-c1">0</span>)
    <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>;

  <span class="pl-smi">int</span> <span class="pl-s1">y_low</span> <span class="pl-c1">=</span> (<span class="pl-smi">int</span>)<span class="pl-s1">y</span>;
  <span class="pl-smi">int</span> <span class="pl-s1">x_low</span> <span class="pl-c1">=</span> (<span class="pl-smi">int</span>)<span class="pl-s1">x</span>;
  <span class="pl-smi">int</span> <span class="pl-s1">y_high</span>;
  <span class="pl-smi">int</span> <span class="pl-s1">x_high</span>;

  <span class="pl-k">if</span> (<span class="pl-s1">y_low</span> &gt;= <span class="pl-s1">height</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>) {
    <span class="pl-s1">y_high</span> <span class="pl-c1">=</span> <span class="pl-s1">y_low</span> <span class="pl-c1">=</span> <span class="pl-s1">height</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>;
    <span class="pl-s1">y</span> <span class="pl-c1">=</span> (<span class="pl-smi">T</span>)<span class="pl-s1">y_low</span>;
  } <span class="pl-k">else</span> {
    <span class="pl-s1">y_high</span> <span class="pl-c1">=</span> <span class="pl-s1">y_low</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>;
  }

  <span class="pl-k">if</span> (<span class="pl-s1">x_low</span> &gt;= <span class="pl-s1">width</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>) {
    <span class="pl-s1">x_high</span> <span class="pl-c1">=</span> <span class="pl-s1">x_low</span> <span class="pl-c1">=</span> <span class="pl-s1">width</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>;
    <span class="pl-s1">x</span> <span class="pl-c1">=</span> (<span class="pl-smi">T</span>)<span class="pl-s1">x_low</span>;
  } <span class="pl-k">else</span> {
    <span class="pl-s1">x_high</span> <span class="pl-c1">=</span> <span class="pl-s1">x_low</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>;
  }

  <span class="pl-smi">T</span> <span class="pl-s1">ly</span> <span class="pl-c1">=</span> <span class="pl-s1">y</span> <span class="pl-c1">-</span> <span class="pl-s1">y_low</span>;
  <span class="pl-smi">T</span> <span class="pl-s1">lx</span> <span class="pl-c1">=</span> <span class="pl-s1">x</span> <span class="pl-c1">-</span> <span class="pl-s1">x_low</span>;
  <span class="pl-smi">T</span> <span class="pl-s1">hy</span> <span class="pl-c1">=</span> <span class="pl-c1">1.</span> <span class="pl-c1">-</span> <span class="pl-s1">ly</span>, <span class="pl-s1">hx</span> <span class="pl-c1">=</span> <span class="pl-c1">1.</span> <span class="pl-c1">-</span> <span class="pl-s1">lx</span>;

  <span class="pl-c">// do bilinear interpolation</span>
  <span class="pl-smi">T</span> <span class="pl-s1">v1</span> <span class="pl-c1">=</span> <span class="pl-s1">input</span>[<span class="pl-s1">y_low</span> <span class="pl-c1">*</span> <span class="pl-s1">width</span> <span class="pl-c1">+</span> <span class="pl-s1">x_low</span>];
  <span class="pl-smi">T</span> <span class="pl-s1">v2</span> <span class="pl-c1">=</span> <span class="pl-s1">input</span>[<span class="pl-s1">y_low</span> <span class="pl-c1">*</span> <span class="pl-s1">width</span> <span class="pl-c1">+</span> <span class="pl-s1">x_high</span>];
  <span class="pl-smi">T</span> <span class="pl-s1">v3</span> <span class="pl-c1">=</span> <span class="pl-s1">input</span>[<span class="pl-s1">y_high</span> <span class="pl-c1">*</span> <span class="pl-s1">width</span> <span class="pl-c1">+</span> <span class="pl-s1">x_low</span>];
  <span class="pl-smi">T</span> <span class="pl-s1">v4</span> <span class="pl-c1">=</span> <span class="pl-s1">input</span>[<span class="pl-s1">y_high</span> <span class="pl-c1">*</span> <span class="pl-s1">width</span> <span class="pl-c1">+</span> <span class="pl-s1">x_high</span>];
  <span class="pl-smi">T</span> <span class="pl-s1">w1</span> <span class="pl-c1">=</span> <span class="pl-s1">hy</span> <span class="pl-c1">*</span> <span class="pl-s1">hx</span>, <span class="pl-s1">w2</span> <span class="pl-c1">=</span> <span class="pl-s1">hy</span> <span class="pl-c1">*</span> <span class="pl-s1">lx</span>, <span class="pl-s1">w3</span> <span class="pl-c1">=</span> <span class="pl-s1">ly</span> <span class="pl-c1">*</span> <span class="pl-s1">hx</span>, <span class="pl-s1">w4</span> <span class="pl-c1">=</span> <span class="pl-s1">ly</span> <span class="pl-c1">*</span> <span class="pl-s1">lx</span>;

  <span class="pl-smi">T</span> <span class="pl-s1">val</span> <span class="pl-c1">=</span> (<span class="pl-s1">w1</span> <span class="pl-c1">*</span> <span class="pl-s1">v1</span> <span class="pl-c1">+</span> <span class="pl-s1">w2</span> <span class="pl-c1">*</span> <span class="pl-s1">v2</span> <span class="pl-c1">+</span> <span class="pl-s1">w3</span> <span class="pl-c1">*</span> <span class="pl-s1">v3</span> <span class="pl-c1">+</span> <span class="pl-s1">w4</span> <span class="pl-c1">*</span> <span class="pl-s1">v4</span>);

  <span class="pl-k">return</span> <span class="pl-s1">val</span>;
}

<span class="pl-s1">template</span> <span class="pl-c1">&lt;</span><span class="pl-s1">typename</span> <span class="pl-c1">T</span><span class="pl-c1">&gt;</span>
<span class="pl-s1">__global__</span> <span class="pl-smi">void</span> <span class="pl-en">PSROIAlignForwardCUDA</span>(
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">nthreads</span>,
    <span class="pl-k">const</span> <span class="pl-smi">T</span><span class="pl-c1">*</span> <span class="pl-s1">input</span>,
    <span class="pl-k">const</span> <span class="pl-smi">T</span> <span class="pl-s1">spatial_scale</span>,
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">channels</span>,
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">height</span>,
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">width</span>,
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">pooled_height</span>,
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">pooled_width</span>,
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">sampling_ratio</span>,
    <span class="pl-k">const</span> <span class="pl-smi">T</span><span class="pl-c1">*</span> <span class="pl-s1">rois</span>,
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">channels_out</span>,
    <span class="pl-smi">T</span><span class="pl-c1">*</span> <span class="pl-s1">output</span>,
    <span class="pl-smi">int</span><span class="pl-c1">*</span> <span class="pl-s1">channel_mapping</span>) {
  <span class="pl-en">CUDA_1D_KERNEL_LOOP</span>(<span class="pl-s1">index</span>, <span class="pl-s1">nthreads</span>) {
    <span class="pl-c">// (n, c_out, ph, pw) is an element in the pooled output</span>
    <span class="pl-smi">int</span> <span class="pl-s1">pw</span> <span class="pl-c1">=</span> <span class="pl-s1">index</span> % <span class="pl-s1">pooled_width</span>;
    <span class="pl-smi">int</span> <span class="pl-s1">ph</span> <span class="pl-c1">=</span> (<span class="pl-s1">index</span> / <span class="pl-s1">pooled_width</span>) % <span class="pl-s1">pooled_height</span>;
    <span class="pl-smi">int</span> <span class="pl-s1">c_out</span> <span class="pl-c1">=</span> (<span class="pl-s1">index</span> / <span class="pl-s1">pooled_width</span> / <span class="pl-s1">pooled_height</span>) % <span class="pl-s1">channels_out</span>;
    <span class="pl-smi">int</span> <span class="pl-s1">n</span> <span class="pl-c1">=</span> <span class="pl-s1">index</span> / <span class="pl-s1">pooled_width</span> / <span class="pl-s1">pooled_height</span> / <span class="pl-s1">channels_out</span>;

    <span class="pl-c">// (n, c_in, ph, pw) is the associated element in the input</span>
    <span class="pl-smi">int</span> <span class="pl-s1">c_in</span> <span class="pl-c1">=</span> (<span class="pl-s1">c_out</span> <span class="pl-c1">*</span> <span class="pl-s1">pooled_height</span> <span class="pl-c1">+</span> <span class="pl-s1">ph</span>) <span class="pl-c1">*</span> <span class="pl-s1">pooled_width</span> <span class="pl-c1">+</span> <span class="pl-s1">pw</span>;

    <span class="pl-c">// [start, end) interval for spatial sampling</span>
    <span class="pl-k">const</span> <span class="pl-smi">T</span><span class="pl-c1">*</span> <span class="pl-s1">offset_rois</span> <span class="pl-c1">=</span> <span class="pl-s1">rois</span> <span class="pl-c1">+</span> <span class="pl-s1">n</span> <span class="pl-c1">*</span> <span class="pl-c1">5</span>;
    <span class="pl-smi">int</span> <span class="pl-s1">roi_batch_ind</span> <span class="pl-c1">=</span> <span class="pl-s1">offset_rois</span>[<span class="pl-c1">0</span>];

    <span class="pl-c">// Do not using rounding; this implementation detail is critical</span>
    <span class="pl-smi">T</span> <span class="pl-s1">roi_start_w</span> <span class="pl-c1">=</span> <span class="pl-s1">offset_rois</span>[<span class="pl-c1">1</span>] <span class="pl-c1">*</span> <span class="pl-s1">spatial_scale</span> <span class="pl-c1">-</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-c1">0.5</span>);
    <span class="pl-smi">T</span> <span class="pl-s1">roi_start_h</span> <span class="pl-c1">=</span> <span class="pl-s1">offset_rois</span>[<span class="pl-c1">2</span>] <span class="pl-c1">*</span> <span class="pl-s1">spatial_scale</span> <span class="pl-c1">-</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-c1">0.5</span>);
    <span class="pl-smi">T</span> <span class="pl-s1">roi_end_w</span> <span class="pl-c1">=</span> <span class="pl-s1">offset_rois</span>[<span class="pl-c1">3</span>] <span class="pl-c1">*</span> <span class="pl-s1">spatial_scale</span> <span class="pl-c1">-</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-c1">0.5</span>);
    <span class="pl-smi">T</span> <span class="pl-s1">roi_end_h</span> <span class="pl-c1">=</span> <span class="pl-s1">offset_rois</span>[<span class="pl-c1">4</span>] <span class="pl-c1">*</span> <span class="pl-s1">spatial_scale</span> <span class="pl-c1">-</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-c1">0.5</span>);

    <span class="pl-smi">T</span> <span class="pl-s1">roi_width</span> <span class="pl-c1">=</span> <span class="pl-s1">roi_end_w</span> <span class="pl-c1">-</span> <span class="pl-s1">roi_start_w</span>;
    <span class="pl-smi">T</span> <span class="pl-s1">roi_height</span> <span class="pl-c1">=</span> <span class="pl-s1">roi_end_h</span> <span class="pl-c1">-</span> <span class="pl-s1">roi_start_h</span>;
    <span class="pl-smi">T</span> <span class="pl-s1">bin_size_h</span> <span class="pl-c1">=</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">roi_height</span>) / <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">pooled_height</span>);
    <span class="pl-smi">T</span> <span class="pl-s1">bin_size_w</span> <span class="pl-c1">=</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">roi_width</span>) / <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">pooled_width</span>);

    <span class="pl-c">// Do not using floor/ceil; this implementation detail is critical</span>
    <span class="pl-smi">T</span> <span class="pl-s1">hstart</span> <span class="pl-c1">=</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">ph</span>) <span class="pl-c1">*</span> <span class="pl-s1">bin_size_h</span> <span class="pl-c1">+</span> <span class="pl-s1">roi_start_h</span>;
    <span class="pl-smi">T</span> <span class="pl-s1">wstart</span> <span class="pl-c1">=</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">pw</span>) <span class="pl-c1">*</span> <span class="pl-s1">bin_size_w</span> <span class="pl-c1">+</span> <span class="pl-s1">roi_start_w</span>;

    <span class="pl-c">// We use roi_bin_grid to sample the grid and mimic integral</span>
    <span class="pl-smi">int</span> <span class="pl-s1">roi_bin_grid_h</span> <span class="pl-c1">=</span> (<span class="pl-s1">sampling_ratio</span> <span class="pl-c1">&gt;</span> <span class="pl-c1">0</span>)
        ? <span class="pl-s1">sampling_ratio</span>
        : <span class="pl-en">ceil</span>(<span class="pl-s1">roi_height</span> / <span class="pl-s1">pooled_height</span>);
    <span class="pl-smi">int</span> <span class="pl-s1">roi_bin_grid_w</span> <span class="pl-c1">=</span>
        (<span class="pl-s1">sampling_ratio</span> <span class="pl-c1">&gt;</span> <span class="pl-c1">0</span>) ? <span class="pl-s1">sampling_ratio</span> : <span class="pl-en">ceil</span>(<span class="pl-s1">roi_width</span> / <span class="pl-s1">pooled_width</span>);
    <span class="pl-k">const</span> <span class="pl-smi">T</span> <span class="pl-s1">count</span> <span class="pl-c1">=</span> <span class="pl-s1">roi_bin_grid_h</span> <span class="pl-c1">*</span> <span class="pl-s1">roi_bin_grid_w</span>;

    <span class="pl-k">const</span> <span class="pl-smi">T</span><span class="pl-c1">*</span> <span class="pl-s1">offset_input</span> <span class="pl-c1">=</span>
        <span class="pl-s1">input</span> <span class="pl-c1">+</span> (<span class="pl-s1">roi_batch_ind</span> <span class="pl-c1">*</span> <span class="pl-s1">channels</span> <span class="pl-c1">+</span> <span class="pl-s1">c_in</span>) <span class="pl-c1">*</span> <span class="pl-s1">height</span> <span class="pl-c1">*</span> <span class="pl-s1">width</span>;
    <span class="pl-smi">T</span> <span class="pl-s1">out_sum</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>;
    <span class="pl-k">for</span> (<span class="pl-smi">int</span> <span class="pl-s1">iy</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>; <span class="pl-s1">iy</span> <span class="pl-c1">&lt;</span> <span class="pl-s1">roi_bin_grid_h</span>; <span class="pl-s1">iy</span><span class="pl-c1">++</span>) {
      <span class="pl-k">const</span> <span class="pl-smi">T</span> <span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">hstart</span> <span class="pl-c1">+</span>
          <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">iy</span> <span class="pl-c1">+</span> <span class="pl-c1">.5f</span>) <span class="pl-c1">*</span> <span class="pl-s1">bin_size_h</span> /
              <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">roi_bin_grid_h</span>);
      <span class="pl-k">for</span> (<span class="pl-smi">int</span> <span class="pl-s1">ix</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>; <span class="pl-s1">ix</span> <span class="pl-c1">&lt;</span> <span class="pl-s1">roi_bin_grid_w</span>; <span class="pl-s1">ix</span><span class="pl-c1">++</span>) {
        <span class="pl-k">const</span> <span class="pl-smi">T</span> <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">wstart</span> <span class="pl-c1">+</span>
            <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">ix</span> <span class="pl-c1">+</span> <span class="pl-c1">.5f</span>) <span class="pl-c1">*</span> <span class="pl-s1">bin_size_w</span> /
                <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">roi_bin_grid_w</span>);
        <span class="pl-smi">T</span> <span class="pl-s1">val</span> <span class="pl-c1">=</span> <span class="pl-en">bilinear_interpolate</span>(<span class="pl-s1">offset_input</span>, <span class="pl-s1">height</span>, <span class="pl-s1">width</span>, <span class="pl-s1">y</span>, <span class="pl-s1">x</span>, <span class="pl-s1">index</span>);
        <span class="pl-s1">out_sum</span> <span class="pl-c1">+=</span> <span class="pl-s1">val</span>;
      }
    }

    <span class="pl-s1">out_sum</span> /= <span class="pl-s1">count</span>;
    <span class="pl-s1">output</span>[<span class="pl-s1">index</span>] <span class="pl-c1">=</span> <span class="pl-s1">out_sum</span>;
    <span class="pl-s1">channel_mapping</span>[<span class="pl-s1">index</span>] <span class="pl-c1">=</span> <span class="pl-s1">c_in</span>;
  }
}

<span class="pl-s1">template</span> <span class="pl-c1">&lt;</span><span class="pl-s1">typename</span> <span class="pl-c1">T</span><span class="pl-c1">&gt;</span>
<span class="pl-s1">__device__</span> <span class="pl-smi">void</span> <span class="pl-en">bilinear_interpolate_gradient</span>(
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">height</span>,
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">width</span>,
    <span class="pl-smi">T</span> <span class="pl-s1">y</span>,
    <span class="pl-smi">T</span> <span class="pl-s1">x</span>,
    <span class="pl-smi">T</span><span class="pl-c1">&amp;</span> <span class="pl-s1">w1</span>,
    <span class="pl-smi">T</span><span class="pl-c1">&amp;</span> <span class="pl-s1">w2</span>,
    <span class="pl-smi">T</span><span class="pl-c1">&amp;</span> <span class="pl-s1">w3</span>,
    <span class="pl-smi">T</span><span class="pl-c1">&amp;</span> <span class="pl-s1">w4</span>,
    <span class="pl-smi">int</span><span class="pl-c1">&amp;</span> <span class="pl-s1">x_low</span>,
    <span class="pl-smi">int</span><span class="pl-c1">&amp;</span> <span class="pl-s1">x_high</span>,
    <span class="pl-smi">int</span><span class="pl-c1">&amp;</span> <span class="pl-s1">y_low</span>,
    <span class="pl-smi">int</span><span class="pl-c1">&amp;</span> <span class="pl-s1">y_high</span>,
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">index</span> <span class="pl-c">/* index for debug only*/</span>) {
  <span class="pl-c">// deal with cases that inverse elements are out of feature map boundary</span>
  <span class="pl-k">if</span> (<span class="pl-s1">y</span> <span class="pl-c1">&lt;</span> <span class="pl-c1">-1.0</span> <span class="pl-c1">||</span> <span class="pl-s1">y</span> <span class="pl-c1">&gt;</span> <span class="pl-s1">height</span> <span class="pl-c1">||</span> <span class="pl-s1">x</span> <span class="pl-c1">&lt;</span> <span class="pl-c1">-1.0</span> <span class="pl-c1">||</span> <span class="pl-s1">x</span> <span class="pl-c1">&gt;</span> <span class="pl-s1">width</span>) {
    <span class="pl-c">// empty</span>
    <span class="pl-s1">w1</span> <span class="pl-c1">=</span> <span class="pl-s1">w2</span> <span class="pl-c1">=</span> <span class="pl-s1">w3</span> <span class="pl-c1">=</span> <span class="pl-s1">w4</span> <span class="pl-c1">=</span> <span class="pl-c1">0.</span>;
    <span class="pl-s1">x_low</span> <span class="pl-c1">=</span> <span class="pl-s1">x_high</span> <span class="pl-c1">=</span> <span class="pl-s1">y_low</span> <span class="pl-c1">=</span> <span class="pl-s1">y_high</span> <span class="pl-c1">=</span> <span class="pl-c1">-1</span>;
    <span class="pl-k">return</span>;
  }

  <span class="pl-k">if</span> (<span class="pl-s1">y</span> &lt;= <span class="pl-c1">0</span>)
    <span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>;
  <span class="pl-k">if</span> (<span class="pl-s1">x</span> &lt;= <span class="pl-c1">0</span>)
    <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>;

  <span class="pl-s1">y_low</span> <span class="pl-c1">=</span> (<span class="pl-smi">int</span>)<span class="pl-s1">y</span>;
  <span class="pl-s1">x_low</span> <span class="pl-c1">=</span> (<span class="pl-smi">int</span>)<span class="pl-s1">x</span>;

  <span class="pl-k">if</span> (<span class="pl-s1">y_low</span> &gt;= <span class="pl-s1">height</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>) {
    <span class="pl-s1">y_high</span> <span class="pl-c1">=</span> <span class="pl-s1">y_low</span> <span class="pl-c1">=</span> <span class="pl-s1">height</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>;
    <span class="pl-s1">y</span> <span class="pl-c1">=</span> (<span class="pl-smi">T</span>)<span class="pl-s1">y_low</span>;
  } <span class="pl-k">else</span> {
    <span class="pl-s1">y_high</span> <span class="pl-c1">=</span> <span class="pl-s1">y_low</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>;
  }

  <span class="pl-k">if</span> (<span class="pl-s1">x_low</span> &gt;= <span class="pl-s1">width</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>) {
    <span class="pl-s1">x_high</span> <span class="pl-c1">=</span> <span class="pl-s1">x_low</span> <span class="pl-c1">=</span> <span class="pl-s1">width</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>;
    <span class="pl-s1">x</span> <span class="pl-c1">=</span> (<span class="pl-smi">T</span>)<span class="pl-s1">x_low</span>;
  } <span class="pl-k">else</span> {
    <span class="pl-s1">x_high</span> <span class="pl-c1">=</span> <span class="pl-s1">x_low</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>;
  }

  <span class="pl-smi">T</span> <span class="pl-s1">ly</span> <span class="pl-c1">=</span> <span class="pl-s1">y</span> <span class="pl-c1">-</span> <span class="pl-s1">y_low</span>;
  <span class="pl-smi">T</span> <span class="pl-s1">lx</span> <span class="pl-c1">=</span> <span class="pl-s1">x</span> <span class="pl-c1">-</span> <span class="pl-s1">x_low</span>;
  <span class="pl-smi">T</span> <span class="pl-s1">hy</span> <span class="pl-c1">=</span> <span class="pl-c1">1.</span> <span class="pl-c1">-</span> <span class="pl-s1">ly</span>, <span class="pl-s1">hx</span> <span class="pl-c1">=</span> <span class="pl-c1">1.</span> <span class="pl-c1">-</span> <span class="pl-s1">lx</span>;

  <span class="pl-c">// reference in forward</span>
  <span class="pl-c">// T v1 = input[y_low * width + x_low];</span>
  <span class="pl-c">// T v2 = input[y_low * width + x_high];</span>
  <span class="pl-c">// T v3 = input[y_high * width + x_low];</span>
  <span class="pl-c">// T v4 = input[y_high * width + x_high];</span>
  <span class="pl-c">// T val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);</span>

  <span class="pl-s1">w1</span> <span class="pl-c1">=</span> <span class="pl-s1">hy</span> <span class="pl-c1">*</span> <span class="pl-s1">hx</span>, <span class="pl-s1">w2</span> <span class="pl-c1">=</span> <span class="pl-s1">hy</span> <span class="pl-c1">*</span> <span class="pl-s1">lx</span>, <span class="pl-s1">w3</span> <span class="pl-c1">=</span> <span class="pl-s1">ly</span> <span class="pl-c1">*</span> <span class="pl-s1">hx</span>, <span class="pl-s1">w4</span> <span class="pl-c1">=</span> <span class="pl-s1">ly</span> <span class="pl-c1">*</span> <span class="pl-s1">lx</span>;

  <span class="pl-k">return</span>;
}

<span class="pl-s1">template</span> <span class="pl-c1">&lt;</span><span class="pl-s1">typename</span> <span class="pl-c1">T</span><span class="pl-c1">&gt;</span>
<span class="pl-s1">__global__</span> <span class="pl-smi">void</span> <span class="pl-en">PSROIAlignBackwardCUDA</span>(
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">nthreads</span>,
    <span class="pl-k">const</span> <span class="pl-smi">T</span><span class="pl-c1">*</span> <span class="pl-s1">grad_output</span>,
    <span class="pl-k">const</span> <span class="pl-smi">int</span><span class="pl-c1">*</span> <span class="pl-s1">channel_mapping</span>,
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">num_rois</span>,
    <span class="pl-k">const</span> <span class="pl-smi">T</span> <span class="pl-s1">spatial_scale</span>,
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">channels</span>,
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">height</span>,
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">width</span>,
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">pooled_height</span>,
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">pooled_width</span>,
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">sampling_ratio</span>,
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">channels_out</span>,
    <span class="pl-smi">T</span><span class="pl-c1">*</span> <span class="pl-s1">grad_input</span>,
    <span class="pl-k">const</span> <span class="pl-smi">T</span><span class="pl-c1">*</span> <span class="pl-s1">rois</span>) {
  <span class="pl-en">CUDA_1D_KERNEL_LOOP</span>(<span class="pl-s1">index</span>, <span class="pl-s1">nthreads</span>) {
    <span class="pl-c">// (n, *, ph, pw) is an element in the pooled output</span>
    <span class="pl-smi">int</span> <span class="pl-s1">pw</span> <span class="pl-c1">=</span> <span class="pl-s1">index</span> % <span class="pl-s1">pooled_width</span>;
    <span class="pl-smi">int</span> <span class="pl-s1">ph</span> <span class="pl-c1">=</span> (<span class="pl-s1">index</span> / <span class="pl-s1">pooled_width</span>) % <span class="pl-s1">pooled_height</span>;
    <span class="pl-smi">int</span> <span class="pl-s1">n</span> <span class="pl-c1">=</span> <span class="pl-s1">index</span> / <span class="pl-s1">pooled_width</span> / <span class="pl-s1">pooled_height</span> / <span class="pl-s1">channels_out</span>;

    <span class="pl-k">const</span> <span class="pl-smi">T</span><span class="pl-c1">*</span> <span class="pl-s1">offset_rois</span> <span class="pl-c1">=</span> <span class="pl-s1">rois</span> <span class="pl-c1">+</span> <span class="pl-s1">n</span> <span class="pl-c1">*</span> <span class="pl-c1">5</span>;
    <span class="pl-smi">int</span> <span class="pl-s1">roi_batch_ind</span> <span class="pl-c1">=</span> <span class="pl-s1">offset_rois</span>[<span class="pl-c1">0</span>];

    <span class="pl-c">// Do not using rounding; this implementation detail is critical</span>
    <span class="pl-smi">T</span> <span class="pl-s1">roi_start_w</span> <span class="pl-c1">=</span> <span class="pl-s1">offset_rois</span>[<span class="pl-c1">1</span>] <span class="pl-c1">*</span> <span class="pl-s1">spatial_scale</span> <span class="pl-c1">-</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-c1">0.5</span>);
    <span class="pl-smi">T</span> <span class="pl-s1">roi_start_h</span> <span class="pl-c1">=</span> <span class="pl-s1">offset_rois</span>[<span class="pl-c1">2</span>] <span class="pl-c1">*</span> <span class="pl-s1">spatial_scale</span> <span class="pl-c1">-</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-c1">0.5</span>);
    <span class="pl-smi">T</span> <span class="pl-s1">roi_end_w</span> <span class="pl-c1">=</span> <span class="pl-s1">offset_rois</span>[<span class="pl-c1">3</span>] <span class="pl-c1">*</span> <span class="pl-s1">spatial_scale</span> <span class="pl-c1">-</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-c1">0.5</span>);
    <span class="pl-smi">T</span> <span class="pl-s1">roi_end_h</span> <span class="pl-c1">=</span> <span class="pl-s1">offset_rois</span>[<span class="pl-c1">4</span>] <span class="pl-c1">*</span> <span class="pl-s1">spatial_scale</span> <span class="pl-c1">-</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-c1">0.5</span>);

    <span class="pl-c">// Force too small ROIs to be 1x1</span>
    <span class="pl-smi">T</span> <span class="pl-s1">roi_width</span> <span class="pl-c1">=</span> <span class="pl-s1">roi_end_w</span> <span class="pl-c1">-</span> <span class="pl-s1">roi_start_w</span>;
    <span class="pl-smi">T</span> <span class="pl-s1">roi_height</span> <span class="pl-c1">=</span> <span class="pl-s1">roi_end_h</span> <span class="pl-c1">-</span> <span class="pl-s1">roi_start_h</span>;
    <span class="pl-smi">T</span> <span class="pl-s1">bin_size_h</span> <span class="pl-c1">=</span> <span class="pl-s1">roi_height</span> / <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">pooled_height</span>);
    <span class="pl-smi">T</span> <span class="pl-s1">bin_size_w</span> <span class="pl-c1">=</span> <span class="pl-s1">roi_width</span> / <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">pooled_width</span>);

    <span class="pl-smi">int</span> <span class="pl-s1">c_in</span> <span class="pl-c1">=</span> <span class="pl-s1">channel_mapping</span>[<span class="pl-s1">index</span>];
    <span class="pl-smi">T</span><span class="pl-c1">*</span> <span class="pl-s1">grad_input_offset</span> <span class="pl-c1">=</span>
        <span class="pl-s1">grad_input</span> <span class="pl-c1">+</span> (<span class="pl-s1">roi_batch_ind</span> <span class="pl-c1">*</span> <span class="pl-s1">channels</span> <span class="pl-c1">+</span> <span class="pl-s1">c_in</span>) <span class="pl-c1">*</span> <span class="pl-s1">height</span> <span class="pl-c1">*</span> <span class="pl-s1">width</span>;

    <span class="pl-c">// Do not using floor/ceil; this implementation detail is critical</span>
    <span class="pl-smi">T</span> <span class="pl-s1">hstart</span> <span class="pl-c1">=</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">ph</span>) <span class="pl-c1">*</span> <span class="pl-s1">bin_size_h</span> <span class="pl-c1">+</span> <span class="pl-s1">roi_start_h</span>;
    <span class="pl-smi">T</span> <span class="pl-s1">wstart</span> <span class="pl-c1">=</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">pw</span>) <span class="pl-c1">*</span> <span class="pl-s1">bin_size_w</span> <span class="pl-c1">+</span> <span class="pl-s1">roi_start_w</span>;

    <span class="pl-k">const</span> <span class="pl-smi">T</span> <span class="pl-s1">grad_output_this_bin</span> <span class="pl-c1">=</span> <span class="pl-s1">grad_output</span>[<span class="pl-s1">index</span>];

    <span class="pl-c">// We use roi_bin_grid to sample the grid and mimic integral</span>
    <span class="pl-smi">int</span> <span class="pl-s1">roi_bin_grid_h</span> <span class="pl-c1">=</span> (<span class="pl-s1">sampling_ratio</span> <span class="pl-c1">&gt;</span> <span class="pl-c1">0</span>)
        ? <span class="pl-s1">sampling_ratio</span>
        : <span class="pl-en">ceil</span>(<span class="pl-s1">roi_height</span> / <span class="pl-s1">pooled_height</span>); <span class="pl-c">// e.g., = 2</span>
    <span class="pl-smi">int</span> <span class="pl-s1">roi_bin_grid_w</span> <span class="pl-c1">=</span>
        (<span class="pl-s1">sampling_ratio</span> <span class="pl-c1">&gt;</span> <span class="pl-c1">0</span>) ? <span class="pl-s1">sampling_ratio</span> : <span class="pl-en">ceil</span>(<span class="pl-s1">roi_width</span> / <span class="pl-s1">pooled_width</span>);
    <span class="pl-k">const</span> <span class="pl-smi">T</span> <span class="pl-s1">count</span> <span class="pl-c1">=</span> <span class="pl-s1">roi_bin_grid_h</span> <span class="pl-c1">*</span> <span class="pl-s1">roi_bin_grid_w</span>;

    <span class="pl-k">for</span> (<span class="pl-smi">int</span> <span class="pl-s1">iy</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>; <span class="pl-s1">iy</span> <span class="pl-c1">&lt;</span> <span class="pl-s1">roi_bin_grid_h</span>; <span class="pl-s1">iy</span><span class="pl-c1">++</span>) {
      <span class="pl-k">const</span> <span class="pl-smi">T</span> <span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">hstart</span> <span class="pl-c1">+</span>
          <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">iy</span> <span class="pl-c1">+</span> <span class="pl-c1">.5f</span>) <span class="pl-c1">*</span> <span class="pl-s1">bin_size_h</span> /
              <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">roi_bin_grid_h</span>);
      <span class="pl-k">for</span> (<span class="pl-smi">int</span> <span class="pl-s1">ix</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>; <span class="pl-s1">ix</span> <span class="pl-c1">&lt;</span> <span class="pl-s1">roi_bin_grid_w</span>; <span class="pl-s1">ix</span><span class="pl-c1">++</span>) {
        <span class="pl-k">const</span> <span class="pl-smi">T</span> <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">wstart</span> <span class="pl-c1">+</span>
            <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">ix</span> <span class="pl-c1">+</span> <span class="pl-c1">.5f</span>) <span class="pl-c1">*</span> <span class="pl-s1">bin_size_w</span> /
                <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">roi_bin_grid_w</span>);

        <span class="pl-smi">T</span> <span class="pl-s1">w1</span>, <span class="pl-s1">w2</span>, <span class="pl-s1">w3</span>, <span class="pl-s1">w4</span>;
        <span class="pl-smi">int</span> <span class="pl-s1">x_low</span>, <span class="pl-s1">x_high</span>, <span class="pl-s1">y_low</span>, <span class="pl-s1">y_high</span>;

        <span class="pl-en">bilinear_interpolate_gradient</span>(
            <span class="pl-s1">height</span>,
            <span class="pl-s1">width</span>,
            <span class="pl-s1">y</span>,
            <span class="pl-s1">x</span>,
            <span class="pl-s1">w1</span>,
            <span class="pl-s1">w2</span>,
            <span class="pl-s1">w3</span>,
            <span class="pl-s1">w4</span>,
            <span class="pl-s1">x_low</span>,
            <span class="pl-s1">x_high</span>,
            <span class="pl-s1">y_low</span>,
            <span class="pl-s1">y_high</span>,
            <span class="pl-s1">index</span>);

        <span class="pl-smi">T</span> <span class="pl-s1">g1</span> <span class="pl-c1">=</span> <span class="pl-s1">grad_output_this_bin</span> <span class="pl-c1">*</span> <span class="pl-s1">w1</span> / <span class="pl-s1">count</span>;
        <span class="pl-smi">T</span> <span class="pl-s1">g2</span> <span class="pl-c1">=</span> <span class="pl-s1">grad_output_this_bin</span> <span class="pl-c1">*</span> <span class="pl-s1">w2</span> / <span class="pl-s1">count</span>;
        <span class="pl-smi">T</span> <span class="pl-s1">g3</span> <span class="pl-c1">=</span> <span class="pl-s1">grad_output_this_bin</span> <span class="pl-c1">*</span> <span class="pl-s1">w3</span> / <span class="pl-s1">count</span>;
        <span class="pl-smi">T</span> <span class="pl-s1">g4</span> <span class="pl-c1">=</span> <span class="pl-s1">grad_output_this_bin</span> <span class="pl-c1">*</span> <span class="pl-s1">w4</span> / <span class="pl-s1">count</span>;

        <span class="pl-k">if</span> (<span class="pl-s1">x_low</span> &gt;= <span class="pl-c1">0</span> <span class="pl-c1">&amp;&amp;</span> <span class="pl-s1">x_high</span> &gt;= <span class="pl-c1">0</span> <span class="pl-c1">&amp;&amp;</span> <span class="pl-s1">y_low</span> &gt;= <span class="pl-c1">0</span> <span class="pl-c1">&amp;&amp;</span> <span class="pl-s1">y_high</span> &gt;= <span class="pl-c1">0</span>) {
          <span class="pl-en">atomicAdd</span>(<span class="pl-s1">grad_input_offset</span> <span class="pl-c1">+</span> <span class="pl-s1">y_low</span> <span class="pl-c1">*</span> <span class="pl-s1">width</span> <span class="pl-c1">+</span> <span class="pl-s1">x_low</span>, <span class="pl-s1">g1</span>);
          <span class="pl-en">atomicAdd</span>(<span class="pl-s1">grad_input_offset</span> <span class="pl-c1">+</span> <span class="pl-s1">y_low</span> <span class="pl-c1">*</span> <span class="pl-s1">width</span> <span class="pl-c1">+</span> <span class="pl-s1">x_high</span>, <span class="pl-s1">g2</span>);
          <span class="pl-en">atomicAdd</span>(<span class="pl-s1">grad_input_offset</span> <span class="pl-c1">+</span> <span class="pl-s1">y_high</span> <span class="pl-c1">*</span> <span class="pl-s1">width</span> <span class="pl-c1">+</span> <span class="pl-s1">x_low</span>, <span class="pl-s1">g3</span>);
          <span class="pl-en">atomicAdd</span>(<span class="pl-s1">grad_input_offset</span> <span class="pl-c1">+</span> <span class="pl-s1">y_high</span> <span class="pl-c1">*</span> <span class="pl-s1">width</span> <span class="pl-c1">+</span> <span class="pl-s1">x_high</span>, <span class="pl-s1">g4</span>);
        } <span class="pl-c">// if</span>
      } <span class="pl-c">// ix</span>
    } <span class="pl-c">// iy</span>
  }
}

<span class="pl-ent">std</span>::<span class="pl-s1">tuple</span><span class="pl-c1">&lt;</span><span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span>, <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&gt;</span> <span class="pl-en">PSROIAlign_forward_cuda</span>(
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">input</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">rois</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">float</span> <span class="pl-s1">spatial_scale</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">pooled_height</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">pooled_width</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">sampling_ratio</span>) {
  <span class="pl-c">// Check if input tensors are CUDA tensors</span>
  <span class="pl-en">AT_ASSERTM</span>(<span class="pl-s1">input</span>.<span class="pl-en">type</span>().<span class="pl-en">is_cuda</span>(), <span class="pl-s">"input must be a CUDA tensor"</span>);
  <span class="pl-en">AT_ASSERTM</span>(<span class="pl-s1">rois</span>.<span class="pl-en">type</span>().<span class="pl-en">is_cuda</span>(), <span class="pl-s">"rois must be a CUDA tensor"</span>);

  <span class="pl-s1">at</span>::<span class="pl-s1">TensorArg</span> <span class="pl-s1">input_t</span>{<span class="pl-s1">input</span>, <span class="pl-s">"input"</span>, <span class="pl-c1">1</span>}, <span class="pl-s1">rois_t</span>{<span class="pl-s1">rois</span>, <span class="pl-s">"rois"</span>, <span class="pl-c1">2</span>};

  <span class="pl-ent">at</span>::<span class="pl-s1">CheckedFrom</span> <span class="pl-s1">c</span> <span class="pl-c1">=</span> <span class="pl-s">"PSROIAlign_forward_cuda"</span>;
  <span class="pl-ent">at</span>::<span class="pl-s1">checkAllSameGPU</span>(<span class="pl-s1">c</span>, {<span class="pl-s1">input_t</span>, <span class="pl-s1">rois_t</span>});
  <span class="pl-ent">at</span>::<span class="pl-s1">checkAllSameType</span>(<span class="pl-s1">c</span>, {<span class="pl-s1">input_t</span>, <span class="pl-s1">rois_t</span>});

  <span class="pl-ent">at</span>::<span class="pl-ent">cuda</span>::<span class="pl-s1">CUDAGuard</span> <span class="pl-en">device_guard</span>(<span class="pl-s1">input</span>.<span class="pl-en">device</span>());

  auto <span class="pl-smi">num_rois</span><span class="pl-s1"></span> <span class="pl-c1">=</span> <span class="pl-s1">rois</span>.<span class="pl-en">size</span>(<span class="pl-c1">0</span>);
  auto <span class="pl-smi">channels</span><span class="pl-s1"></span> <span class="pl-c1">=</span> <span class="pl-s1">input</span>.<span class="pl-en">size</span>(<span class="pl-c1">1</span>);
  auto <span class="pl-smi">height</span><span class="pl-s1"></span> <span class="pl-c1">=</span> <span class="pl-s1">input</span>.<span class="pl-en">size</span>(<span class="pl-c1">2</span>);
  auto <span class="pl-smi">width</span><span class="pl-s1"></span> <span class="pl-c1">=</span> <span class="pl-s1">input</span>.<span class="pl-en">size</span>(<span class="pl-c1">3</span>);

  <span class="pl-en">AT_ASSERTM</span>(
      <span class="pl-s1">channels</span> % (<span class="pl-s1">pooled_height</span> <span class="pl-c1">*</span> <span class="pl-s1">pooled_width</span>) <span class="pl-c1">==</span> <span class="pl-c1">0</span>,
      <span class="pl-s">"input channels must be a multiple of pooling height * pooling width"</span>);
  <span class="pl-smi">int</span> <span class="pl-s1">channels_out</span> <span class="pl-c1">=</span> <span class="pl-s1">channels</span> / (<span class="pl-s1">pooled_height</span> <span class="pl-c1">*</span> <span class="pl-s1">pooled_width</span>);

  auto <span class="pl-smi">output</span> <span class="pl-c1">=</span> <span class="pl-s1">at</span>::<span class="pl-en">zeros</span>(
      {<span class="pl-smi">num_rois</span>, <span class="pl-smi">channels_out</span>, <span class="pl-smi">pooled_height</span>, <span class="pl-smi">pooled_width</span>}, <span class="pl-smi">input</span>.<span class="pl-en">options</span>());
  auto <span class="pl-smi">channel_mapping</span> <span class="pl-c1">=</span>
      <span class="pl-s1">at</span>::<span class="pl-en">zeros</span>(<span class="pl-smi">output</span>.<span class="pl-en">sizes</span>(), <span class="pl-smi">input</span>.<span class="pl-en">options</span>().<span class="pl-en">dtype</span>(<span class="pl-smi">at</span>::<span class="pl-s1">kInt</span>));

  auto <span class="pl-smi">output_size</span><span class="pl-s1"></span> <span class="pl-c1">=</span> <span class="pl-s1">output</span>.<span class="pl-en">numel</span>();
  <span class="pl-k">if</span> (<span class="pl-s1">output_size</span> <span class="pl-c1">==</span> <span class="pl-c1">0</span>) {
    <span class="pl-en">AT_CUDA_CHECK</span>(<span class="pl-en">cudaGetLastError</span>());
    <span class="pl-k">return</span> <span class="pl-s1">std</span>::<span class="pl-en">make_tuple</span>(<span class="pl-s1">output</span>, <span class="pl-s1">channel_mapping</span>);
  }

  <span class="pl-smi">cudaStream_t</span> <span class="pl-s1">stream</span> <span class="pl-c1">=</span> <span class="pl-s1">at</span>::<span class="pl-s1">cuda</span>::<span class="pl-en">getCurrentCUDAStream</span>();

  <span class="pl-smi">dim3</span> <span class="pl-en">grid</span>(<span class="pl-smi">std</span>::<span class="pl-en">min</span>(
      <span class="pl-smi">at</span>::<span class="pl-s1">cuda</span>::<span class="pl-en">ATenCeilDiv</span>(
          <span class="pl-smi">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-smi">int64_t</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">output_size</span>), <span class="pl-smi">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-smi">int64_t</span><span class="pl-c1">&gt;</span>(<span class="pl-c1">512</span>)),
      <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-smi">int64_t</span><span class="pl-c1">&gt;</span>(<span class="pl-c1">4096</span>)));
  <span class="pl-smi">dim3</span> <span class="pl-en">block</span>(<span class="pl-c1">512</span>);

  <span class="pl-en">AT_DISPATCH_FLOATING_TYPES_AND_HALF</span>(
      <span class="pl-s1">input</span>.<span class="pl-en">scalar_type</span>(), <span class="pl-s">"PSROIAlign_forward"</span>, [<span class="pl-c1">&amp;</span>] {
        <span class="pl-s1">PSROIAlignForwardCUDA</span><span class="pl-c1">&lt;</span><span class="pl-s1">scalar_t</span><span class="pl-c1">&gt;</span>&lt;&lt;<span class="pl-c1">&lt;</span><span class="pl-s1">grid</span>, <span class="pl-s1">block</span>, <span class="pl-c1">0</span>, <span class="pl-s1">stream</span>&gt;&gt;<span class="pl-s1"></span><span class="pl-c1">&gt;</span>(
            <span class="pl-s1">output_size</span>,
            <span class="pl-s1">input</span>.<span class="pl-en">contiguous</span>().<span class="pl-c1">data_ptr</span><span class="pl-c1">&lt;</span><span class="pl-s1">scalar_t</span><span class="pl-c1">&gt;</span>(<span class="pl-s1"></span>),
            <span class="pl-s1">spatial_scale</span>,
            <span class="pl-s1">channels</span>,
            <span class="pl-s1">height</span>,
            <span class="pl-s1">width</span>,
            <span class="pl-s1">pooled_height</span>,
            <span class="pl-s1">pooled_width</span>,
            <span class="pl-s1">sampling_ratio</span>,
            <span class="pl-s1">rois</span>.<span class="pl-en">contiguous</span>().<span class="pl-c1">data_ptr</span><span class="pl-c1">&lt;</span><span class="pl-s1">scalar_t</span><span class="pl-c1">&gt;</span>(<span class="pl-s1"></span>),
            <span class="pl-s1">channels_out</span>,
            <span class="pl-s1">output</span>.<span class="pl-c1">data_ptr</span><span class="pl-c1">&lt;</span><span class="pl-s1">scalar_t</span><span class="pl-c1">&gt;</span>(<span class="pl-s1"></span>),
            <span class="pl-s1">channel_mapping</span>.<span class="pl-c1">data_ptr</span><span class="pl-c1">&lt;</span><span class="pl-s1">int</span><span class="pl-c1">&gt;</span>(<span class="pl-s1"></span>));
      });
  <span class="pl-en">AT_CUDA_CHECK</span>(<span class="pl-en">cudaGetLastError</span>());
  <span class="pl-en">cudaDeviceSynchronize</span>();
  <span class="pl-k">return</span> <span class="pl-s1">std</span>::<span class="pl-en">make_tuple</span>(<span class="pl-s1">output</span>, <span class="pl-s1">channel_mapping</span>);
}

<span class="pl-ent">at</span>::<span class="pl-s1">Tensor</span> <span class="pl-en">PSROIAlign_backward_cuda</span>(
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">grad</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">rois</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">channel_mapping</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">float</span> <span class="pl-s1">spatial_scale</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">pooled_height</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">pooled_width</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">sampling_ratio</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">batch_size</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">channels</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">height</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">width</span>) {
  <span class="pl-c">// Check if input tensors are CUDA tensors</span>
  <span class="pl-en">AT_ASSERTM</span>(<span class="pl-s1">grad</span>.<span class="pl-en">type</span>().<span class="pl-en">is_cuda</span>(), <span class="pl-s">"grad must be a CUDA tensor"</span>);
  <span class="pl-en">AT_ASSERTM</span>(<span class="pl-s1">rois</span>.<span class="pl-en">type</span>().<span class="pl-en">is_cuda</span>(), <span class="pl-s">"rois must be a CUDA tensor"</span>);
  <span class="pl-en">AT_ASSERTM</span>(
      <span class="pl-s1">channel_mapping</span>.<span class="pl-en">type</span>().<span class="pl-en">is_cuda</span>(),
      <span class="pl-s">"channel_mapping must be a CUDA tensor"</span>);

  <span class="pl-s1">at</span>::<span class="pl-s1">TensorArg</span> <span class="pl-s1">grad_t</span>{<span class="pl-s1">grad</span>, <span class="pl-s">"grad"</span>, <span class="pl-c1">1</span>}, <span class="pl-s1">rois_t</span>{<span class="pl-s1">rois</span>, <span class="pl-s">"rois"</span>, <span class="pl-c1">2</span>},
      <span class="pl-s1">channel_mapping_t</span>{<span class="pl-s1">channel_mapping</span>, <span class="pl-s">"channel_mapping"</span>, <span class="pl-c1">3</span>};

  <span class="pl-ent">at</span>::<span class="pl-s1">CheckedFrom</span> <span class="pl-s1">c</span> <span class="pl-c1">=</span> <span class="pl-s">"PSROIAlign_backward_cuda"</span>;
  <span class="pl-ent">at</span>::<span class="pl-s1">checkAllSameGPU</span>(<span class="pl-s1">c</span>, {<span class="pl-s1">grad_t</span>, <span class="pl-s1">rois_t</span>, <span class="pl-s1">channel_mapping_t</span>});
  <span class="pl-ent">at</span>::<span class="pl-s1">checkAllSameType</span>(<span class="pl-s1">c</span>, {<span class="pl-s1">grad_t</span>, <span class="pl-s1">rois_t</span>});

  <span class="pl-ent">at</span>::<span class="pl-ent">cuda</span>::<span class="pl-s1">CUDAGuard</span> <span class="pl-en">device_guard</span>(<span class="pl-s1">grad</span>.<span class="pl-en">device</span>());

  auto <span class="pl-smi">num_rois</span><span class="pl-s1"></span> <span class="pl-c1">=</span> <span class="pl-s1">rois</span>.<span class="pl-en">size</span>(<span class="pl-c1">0</span>);
  auto <span class="pl-smi">grad_input</span> <span class="pl-c1">=</span>
      <span class="pl-s1">at</span>::<span class="pl-en">zeros</span>({<span class="pl-smi">batch_size</span>, <span class="pl-smi">channels</span>, <span class="pl-smi">height</span>, <span class="pl-smi">width</span>}, <span class="pl-smi">grad</span>.<span class="pl-en">options</span>());

  <span class="pl-smi">cudaStream_t</span> <span class="pl-s1">stream</span> <span class="pl-c1">=</span> <span class="pl-s1">at</span>::<span class="pl-s1">cuda</span>::<span class="pl-en">getCurrentCUDAStream</span>();

  <span class="pl-smi">dim3</span> <span class="pl-en">grid</span>(<span class="pl-smi">std</span>::<span class="pl-en">min</span>(
      <span class="pl-smi">at</span>::<span class="pl-s1">cuda</span>::<span class="pl-en">ATenCeilDiv</span>(
          <span class="pl-smi">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-smi">int64_t</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">grad</span>.<span class="pl-en">numel</span>()), <span class="pl-smi">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-smi">int64_t</span><span class="pl-c1">&gt;</span>(<span class="pl-c1">512</span>)),
      <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-smi">int64_t</span><span class="pl-c1">&gt;</span>(<span class="pl-c1">4096</span>)));
  <span class="pl-smi">dim3</span> <span class="pl-en">block</span>(<span class="pl-c1">512</span>);

  <span class="pl-c">// handle possibly empty gradients</span>
  <span class="pl-k">if</span> (<span class="pl-s1">grad</span>.<span class="pl-en">numel</span>() <span class="pl-c1">==</span> <span class="pl-c1">0</span>) {
    <span class="pl-en">AT_CUDA_CHECK</span>(<span class="pl-en">cudaGetLastError</span>());
    <span class="pl-k">return</span> <span class="pl-s1">grad_input</span>;
  }

  <span class="pl-smi">int</span> <span class="pl-s1">channels_out</span> <span class="pl-c1">=</span> <span class="pl-s1">channels</span> / (<span class="pl-s1">pooled_height</span> <span class="pl-c1">*</span> <span class="pl-s1">pooled_width</span>);

  <span class="pl-en">AT_DISPATCH_FLOATING_TYPES_AND_HALF</span>(
      <span class="pl-s1">grad</span>.<span class="pl-en">scalar_type</span>(), <span class="pl-s">"PSROIAlign_backward"</span>, [<span class="pl-c1">&amp;</span>] {
        <span class="pl-s1">PSROIAlignBackwardCUDA</span><span class="pl-c1">&lt;</span><span class="pl-s1">scalar_t</span><span class="pl-c1">&gt;</span>&lt;&lt;<span class="pl-c1">&lt;</span><span class="pl-s1">grid</span>, <span class="pl-s1">block</span>, <span class="pl-c1">0</span>, <span class="pl-s1">stream</span>&gt;&gt;<span class="pl-s1"></span><span class="pl-c1">&gt;</span>(
            <span class="pl-s1">grad</span>.<span class="pl-en">numel</span>(),
            <span class="pl-s1">grad</span>.<span class="pl-en">contiguous</span>().<span class="pl-c1">data_ptr</span><span class="pl-c1">&lt;</span><span class="pl-s1">scalar_t</span><span class="pl-c1">&gt;</span>(<span class="pl-s1"></span>),
            <span class="pl-s1">channel_mapping</span>.<span class="pl-c1">data_ptr</span><span class="pl-c1">&lt;</span><span class="pl-s1">int</span><span class="pl-c1">&gt;</span>(<span class="pl-s1"></span>),
            <span class="pl-s1">num_rois</span>,
            <span class="pl-s1">spatial_scale</span>,
            <span class="pl-s1">channels</span>,
            <span class="pl-s1">height</span>,
            <span class="pl-s1">width</span>,
            <span class="pl-s1">pooled_height</span>,
            <span class="pl-s1">pooled_width</span>,
            <span class="pl-s1">sampling_ratio</span>,
            <span class="pl-s1">channels_out</span>,
            <span class="pl-s1">grad_input</span>.<span class="pl-c1">data_ptr</span><span class="pl-c1">&lt;</span><span class="pl-s1">scalar_t</span><span class="pl-c1">&gt;</span>(<span class="pl-s1"></span>),
            <span class="pl-s1">rois</span>.<span class="pl-en">contiguous</span>().<span class="pl-c1">data_ptr</span><span class="pl-c1">&lt;</span><span class="pl-s1">scalar_t</span><span class="pl-c1">&gt;</span>(<span class="pl-s1"></span>));
      });
  <span class="pl-en">AT_CUDA_CHECK</span>(<span class="pl-en">cudaGetLastError</span>());
  <span class="pl-k">return</span> <span class="pl-s1">grad_input</span>;
}</pre></div>
</details>
<details>
<summary> PSROIAlign_cuda.h </summary>
<div class="highlight highlight-source-c notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="#pragma once
#include &lt;torch/torch.h&gt;
#include &lt;c10/cuda/CUDAGuard.h&gt;

#define CUDA_1D_KERNEL_LOOP(i, n)                                \
  for (int i = (blockIdx.x * blockDim.x) + threadIdx.x; i &lt; (n); \
       i += (blockDim.x * gridDim.x))


std::tuple&lt;at::Tensor, at::Tensor&gt; PSROIAlign_forward_cuda(
    const at::Tensor&amp; input,
    const at::Tensor&amp; rois,
    const float spatial_scale,
    const int pooled_height,
    const int pooled_width,
    const int sampling_ratio);

at::Tensor PSROIAlign_backward_cuda(
    const at::Tensor&amp; grad,
    const at::Tensor&amp; rois,
    const at::Tensor&amp; mapping_channel,
    const float spatial_scale,
    const int pooled_height,
    const int pooled_width,
    const int sampling_ratio,
    const int batch_size,
    const int channels,
    const int height,
    const int width);"><pre class="notranslate"><span class="pl-k">#pragma</span> once
<span class="pl-k">#include</span> <span class="pl-s">&lt;torch/torch.h&gt;</span>
<span class="pl-k">#include</span> <span class="pl-s">&lt;c10/cuda/CUDAGuard.h&gt;</span>

<span class="pl-k">#define</span> <span class="pl-en">CUDA_1D_KERNEL_LOOP</span>(<span class="pl-s1">i</span>, <span class="pl-s1">n</span>)                                \
  for (int i = (blockIdx.x * blockDim.x) + threadIdx.x; i &lt; (n); \
       i += (blockDim.x * gridDim.x))


<span class="pl-ent">std</span>::<span class="pl-s1">tuple</span><span class="pl-c1">&lt;</span><span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span>, <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&gt;</span> <span class="pl-en">PSROIAlign_forward_cuda</span>(
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">input</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">rois</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">float</span> <span class="pl-s1">spatial_scale</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">pooled_height</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">pooled_width</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">sampling_ratio</span>);

<span class="pl-ent">at</span>::<span class="pl-s1">Tensor</span> <span class="pl-en">PSROIAlign_backward_cuda</span>(
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">grad</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">rois</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">mapping_channel</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">float</span> <span class="pl-s1">spatial_scale</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">pooled_height</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">pooled_width</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">sampling_ratio</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">batch_size</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">channels</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">height</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">width</span>);</pre></div>
</details>
<details>
  <summary> PSROIAlign_cpu.cpp </summary>
<div class="highlight highlight-source-c notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="#include &quot;PSROIAlign.h&quot;

template &lt;typename T&gt;
T bilinear_interpolate(
    const T* input,
    const int height,
    const int width,
    T y,
    T x,
    const int index /* index for debug only*/) {
  // deal with cases that inverse elements are out of feature map boundary
  if (y &lt; -1.0 || y &gt; height || x &lt; -1.0 || x &gt; width) {
    // empty
    return 0;
  }

  if (y &lt;= 0)
    y = 0;
  if (x &lt;= 0)
    x = 0;

  int y_low = (int)y;
  int x_low = (int)x;
  int y_high;
  int x_high;

  if (y_low &gt;= height - 1) {
    y_high = y_low = height - 1;
    y = (T)y_low;
  } else {
    y_high = y_low + 1;
  }

  if (x_low &gt;= width - 1) {
    x_high = x_low = width - 1;
    x = (T)x_low;
  } else {
    x_high = x_low + 1;
  }

  T ly = y - y_low;
  T lx = x - x_low;
  T hy = 1. - ly, hx = 1. - lx;

  // do bilinear interpolation
  T v1 = input[y_low * width + x_low];
  T v2 = input[y_low * width + x_high];
  T v3 = input[y_high * width + x_low];
  T v4 = input[y_high * width + x_high];
  T w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;

  T val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);

  return val;
}

template &lt;typename T&gt;
void PSROIAlignForwardCPU(
    const int nthreads,
    const T* input,
    const T spatial_scale,
    const int channels,
    const int height,
    const int width,
    const int pooled_height,
    const int pooled_width,
    const int sampling_ratio,
    const T* rois,
    const int channels_out,
    T* output,
    int* channel_mapping) {
  int num_rois = nthreads / channels_out / pooled_width / pooled_height;
  for (int n = 0; n &lt; num_rois; n++) {
    // [start, end) interval for spatial sampling
    const T* offset_rois = rois + n * 5;
    int roi_batch_ind = offset_rois[0];

    // Do not using rounding; this implementation detail is critical
    T roi_start_w = offset_rois[1] * spatial_scale - static_cast&lt;T&gt;(0.5);
    T roi_start_h = offset_rois[2] * spatial_scale - static_cast&lt;T&gt;(0.5);
    T roi_end_w = offset_rois[3] * spatial_scale - static_cast&lt;T&gt;(0.5);
    T roi_end_h = offset_rois[4] * spatial_scale - static_cast&lt;T&gt;(0.5);

    T roi_width = roi_end_w - roi_start_w;
    T roi_height = roi_end_h - roi_start_h;
    T bin_size_h = static_cast&lt;T&gt;(roi_height) / static_cast&lt;T&gt;(pooled_height);
    T bin_size_w = static_cast&lt;T&gt;(roi_width) / static_cast&lt;T&gt;(pooled_width);

    int c_in = 0;
    for (int c_out = 0; c_out &lt; channels_out; ++c_out) {
      for (int ph = 0; ph &lt; pooled_height; ++ph) {
        for (int pw = 0; pw &lt; pooled_width; ++pw) {
          int index =
              ((n * channels_out + c_out) * pooled_height + ph) * pooled_width +
              pw;

          // Do not using floor/ceil; this implementation detail is critical
          T hstart = static_cast&lt;T&gt;(ph) * bin_size_h + roi_start_h;
          T wstart = static_cast&lt;T&gt;(pw) * bin_size_w + roi_start_w;

          // We use roi_bin_grid to sample the grid and mimic integral
          int roi_bin_grid_h = (sampling_ratio &gt; 0)
              ? sampling_ratio
              : ceil(roi_height / pooled_height);
          int roi_bin_grid_w = (sampling_ratio &gt; 0)
              ? sampling_ratio
              : ceil(roi_width / pooled_width);
          const T count = roi_bin_grid_h * roi_bin_grid_w;

          const T* offset_input =
              input + (roi_batch_ind * channels + c_in) * height * width;

          T out_sum = 0;
          for (int iy = 0; iy &lt; roi_bin_grid_h; iy++) {
            const T y = hstart +
                static_cast&lt;T&gt;(iy + .5f) * bin_size_h /
                    static_cast&lt;T&gt;(roi_bin_grid_h);
            for (int ix = 0; ix &lt; roi_bin_grid_w; ix++) {
              const T x = wstart +
                  static_cast&lt;T&gt;(ix + .5f) * bin_size_w /
                      static_cast&lt;T&gt;(roi_bin_grid_w);
              T val = bilinear_interpolate(
                  offset_input, height, width, y, x, index);
              out_sum += val;
            }
          }

          out_sum /= count;
          output[index] = out_sum;
          channel_mapping[index] = c_in;
          c_in++;
        }
      }
    }
  }
}

template &lt;typename T&gt;
void bilinear_interpolate_gradient(
    const int height,
    const int width,
    T y,
    T x,
    T&amp; w1,
    T&amp; w2,
    T&amp; w3,
    T&amp; w4,
    int&amp; x_low,
    int&amp; x_high,
    int&amp; y_low,
    int&amp; y_high,
    const int index /* index for debug only*/) {
  // deal with cases that inverse elements are out of feature map boundary
  if (y &lt; -1.0 || y &gt; height || x &lt; -1.0 || x &gt; width) {
    // empty
    w1 = w2 = w3 = w4 = 0.;
    x_low = x_high = y_low = y_high = -1;
    return;
  }

  if (y &lt;= 0)
    y = 0;
  if (x &lt;= 0)
    x = 0;

  y_low = (int)y;
  x_low = (int)x;

  if (y_low &gt;= height - 1) {
    y_high = y_low = height - 1;
    y = (T)y_low;
  } else {
    y_high = y_low + 1;
  }

  if (x_low &gt;= width - 1) {
    x_high = x_low = width - 1;
    x = (T)x_low;
  } else {
    x_high = x_low + 1;
  }

  T ly = y - y_low;
  T lx = x - x_low;
  T hy = 1. - ly, hx = 1. - lx;

  // reference in forward
  // T v1 = input[y_low * width + x_low];
  // T v2 = input[y_low * width + x_high];
  // T v3 = input[y_high * width + x_low];
  // T v4 = input[y_high * width + x_high];
  // T val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);

  w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;

  return;
}

template &lt;class T&gt;
inline void add(T* address, const T&amp; val) {
  *address += val;
}

template &lt;typename T&gt;
void PSROIAlignBackwardCPU(
    const int nthreads,
    const T* grad_output,
    const int* channel_mapping,
    const int num_rois,
    const T spatial_scale,
    const int channels,
    const int height,
    const int width,
    const int pooled_height,
    const int pooled_width,
    const int sampling_ratio,
    const int channels_out,
    T* grad_input,
    const T* rois) {
  for (int index = 0; index &lt; nthreads; index++) {
    int pw = index % pooled_width;
    int ph = (index / pooled_width) % pooled_height;
    int n = index / pooled_width / pooled_height / channels_out;

    const T* offset_rois = rois + n * 5;
    int roi_batch_ind = offset_rois[0];

    // Do not using rounding; this implementation detail is critical
    T roi_start_w = offset_rois[1] * spatial_scale - static_cast&lt;T&gt;(0.5);
    T roi_start_h = offset_rois[2] * spatial_scale - static_cast&lt;T&gt;(0.5);
    T roi_end_w = offset_rois[3] * spatial_scale - static_cast&lt;T&gt;(0.5);
    T roi_end_h = offset_rois[4] * spatial_scale - static_cast&lt;T&gt;(0.5);

    // Force too small ROIs to be 1x1
    T roi_width = roi_end_w - roi_start_w;
    T roi_height = roi_end_h - roi_start_h;
    T bin_size_h = roi_height / static_cast&lt;T&gt;(pooled_height);
    T bin_size_w = roi_width / static_cast&lt;T&gt;(pooled_width);

    int c_in = channel_mapping[index];
    T* grad_input_offset =
        grad_input + (roi_batch_ind * channels + c_in) * height * width;

    // Do not using floor/ceil; this implementation detail is critical
    T hstart = static_cast&lt;T&gt;(ph) * bin_size_h + roi_start_h;
    T wstart = static_cast&lt;T&gt;(pw) * bin_size_w + roi_start_w;

    const T grad_output_this_bin = grad_output[index];

    // We use roi_bin_grid to sample the grid and mimic integral
    int roi_bin_grid_h = (sampling_ratio &gt; 0)
        ? sampling_ratio
        : ceil(roi_height / pooled_height); // e.g., = 2
    int roi_bin_grid_w =
        (sampling_ratio &gt; 0) ? sampling_ratio : ceil(roi_width / pooled_width);
    const T count = roi_bin_grid_h * roi_bin_grid_w;

    for (int iy = 0; iy &lt; roi_bin_grid_h; iy++) {
      const T y = hstart +
          static_cast&lt;T&gt;(iy + .5f) * bin_size_h /
              static_cast&lt;T&gt;(roi_bin_grid_h);
      for (int ix = 0; ix &lt; roi_bin_grid_w; ix++) {
        const T x = wstart +
            static_cast&lt;T&gt;(ix + .5f) * bin_size_w /
                static_cast&lt;T&gt;(roi_bin_grid_w);

        T w1, w2, w3, w4;
        int x_low, x_high, y_low, y_high;

        bilinear_interpolate_gradient(
            height,
            width,
            y,
            x,
            w1,
            w2,
            w3,
            w4,
            x_low,
            x_high,
            y_low,
            y_high,
            index);

        T g1 = grad_output_this_bin * w1 / count;
        T g2 = grad_output_this_bin * w2 / count;
        T g3 = grad_output_this_bin * w3 / count;
        T g4 = grad_output_this_bin * w4 / count;

        if (x_low &gt;= 0 &amp;&amp; x_high &gt;= 0 &amp;&amp; y_low &gt;= 0 &amp;&amp; y_high &gt;= 0) {
          add(grad_input_offset + y_low * width + x_low, g1);
          add(grad_input_offset + y_low * width + x_high, g2);
          add(grad_input_offset + y_high * width + x_low, g3);
          add(grad_input_offset + y_high * width + x_high, g4);
        } // if
      } // ix
    } // iy
  }
}

std::tuple&lt;at::Tensor, at::Tensor&gt; PSROIAlign_forward_cpu(
    const at::Tensor&amp; input,
    const at::Tensor&amp; rois,
    const float spatial_scale,
    const int pooled_height,
    const int pooled_width,
    const int sampling_ratio) {
  // Check if input tensors are CPU tensors
  AT_ASSERTM(input.device().is_cpu(), &quot;input must be a CPU tensor&quot;);
  AT_ASSERTM(rois.device().is_cpu(), &quot;rois must be a CPU tensor&quot;);

  at::TensorArg input_t{input, &quot;input&quot;, 1}, rois_t{rois, &quot;rois&quot;, 2};

  at::CheckedFrom c = &quot;PSROIAlign_forward_cpu&quot;;
  at::checkAllSameType(c, {input_t, rois_t});

  int num_rois = rois.size(0);
  int channels = input.size(1);
  int height = input.size(2);
  int width = input.size(3);

  AT_ASSERTM(
      channels % (pooled_height * pooled_width) == 0,
      &quot;input channels must be a multiple of pooling height * pooling width&quot;);
  int channels_out = channels / (pooled_height * pooled_width);

  auto output = at::zeros(
      {num_rois, channels_out, pooled_height, pooled_width}, input.options());
  auto channel_mapping =
      at::zeros(output.sizes(), input.options().dtype(at::kInt));

  auto output_size = output.numel();
  if (output_size == 0) {
    return std::make_tuple(output, channel_mapping);
  }

  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
      input.scalar_type(), &quot;PSROIAlign_forward&quot;, [&amp;] {
        PSROIAlignForwardCPU&lt;scalar_t&gt;(
            output_size,
            input.contiguous().data_ptr&lt;scalar_t&gt;(),
            spatial_scale,
            channels,
            height,
            width,
            pooled_height,
            pooled_width,
            sampling_ratio,
            rois.contiguous().data_ptr&lt;scalar_t&gt;(),
            channels_out,
            output.data_ptr&lt;scalar_t&gt;(),
            channel_mapping.data_ptr&lt;int&gt;());
      });
  return std::make_tuple(output, channel_mapping);
}

at::Tensor PSROIAlign_backward_cpu(
    const at::Tensor&amp; grad,
    const at::Tensor&amp; rois,
    const at::Tensor&amp; channel_mapping,
    const float spatial_scale,
    const int pooled_height,
    const int pooled_width,
    const int sampling_ratio,
    const int batch_size,
    const int channels,
    const int height,
    const int width) {
  // Check if input tensors are CPU tensors
  AT_ASSERTM(grad.device().is_cpu(), &quot;grad must be a CPU tensor&quot;);
  AT_ASSERTM(rois.device().is_cpu(), &quot;rois must be a CPU tensor&quot;);
  AT_ASSERTM(
      channel_mapping.device().is_cpu(),
      &quot;channel_mapping must be a CPU tensor&quot;);

  at::TensorArg grad_t{grad, &quot;grad&quot;, 1}, rois_t{rois, &quot;rois&quot;, 2},
      channel_mapping_t{channel_mapping, &quot;channel_mapping&quot;, 3};

  at::CheckedFrom c = &quot;PSROIAlign_backward_cpu&quot;;
  at::checkAllSameType(c, {grad_t, rois_t});

  auto num_rois = rois.size(0);
  auto grad_input =
      at::zeros({batch_size, channels, height, width}, grad.options());

  // handle possibly empty gradients
  if (grad.numel() == 0) {
    return grad_input;
  }

  int channels_out = channels / (pooled_height * pooled_width);

  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
      grad.scalar_type(), &quot;PSROIAlign_backward&quot;, [&amp;] {
        PSROIAlignBackwardCPU&lt;scalar_t&gt;(
            grad.numel(),
            grad.contiguous().data_ptr&lt;scalar_t&gt;(),
            channel_mapping.data_ptr&lt;int&gt;(),
            num_rois,
            spatial_scale,
            channels,
            height,
            width,
            pooled_height,
            pooled_width,
            sampling_ratio,
            channels_out,
            grad_input.data_ptr&lt;scalar_t&gt;(),
            rois.contiguous().data_ptr&lt;scalar_t&gt;());
      });
  return grad_input;
}"><pre class="notranslate"><span class="pl-k">#include</span> <span class="pl-s">"PSROIAlign.h"</span>

<span class="pl-s1">template</span> <span class="pl-c1">&lt;</span><span class="pl-s1">typename</span> <span class="pl-c1">T</span><span class="pl-c1">&gt;</span>
<span class="pl-c1">T</span> <span class="pl-s1">bilinear_interpolate</span>(
    <span class="pl-k">const</span> <span class="pl-smi">T</span><span class="pl-c1">*</span> <span class="pl-s1">input</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">int</span> <span class="pl-s1">height</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">int</span> <span class="pl-s1">width</span>,
    <span class="pl-c1">T</span> <span class="pl-s1">y</span>,
    <span class="pl-c1">T</span> <span class="pl-s1">x</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">int</span> <span class="pl-s1">index</span> <span class="pl-c">/* index for debug only*/</span>) {
  <span class="pl-c">// deal with cases that inverse elements are out of feature map boundary</span>
  <span class="pl-k">if</span> (<span class="pl-s1">y</span> <span class="pl-c1">&lt;</span> <span class="pl-c1">-1.0</span> <span class="pl-c1">||</span> <span class="pl-s1">y</span> <span class="pl-c1">&gt;</span> <span class="pl-s1">height</span> <span class="pl-c1">||</span> <span class="pl-s1">x</span> <span class="pl-c1">&lt;</span> <span class="pl-c1">-1.0</span> <span class="pl-c1">||</span> <span class="pl-s1">x</span> <span class="pl-c1">&gt;</span> <span class="pl-s1">width</span>) {
    <span class="pl-c">// empty</span>
    <span class="pl-k">return</span> <span class="pl-c1">0</span>;
  }

  <span class="pl-k">if</span> (<span class="pl-s1">y</span> &lt;= <span class="pl-c1">0</span>)
    <span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>;
  <span class="pl-k">if</span> (<span class="pl-s1">x</span> &lt;= <span class="pl-c1">0</span>)
    <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>;

  <span class="pl-smi">int</span> <span class="pl-s1">y_low</span> <span class="pl-c1">=</span> (<span class="pl-smi">int</span>)<span class="pl-s1">y</span>;
  <span class="pl-smi">int</span> <span class="pl-s1">x_low</span> <span class="pl-c1">=</span> (<span class="pl-smi">int</span>)<span class="pl-s1">x</span>;
  <span class="pl-smi">int</span> <span class="pl-s1">y_high</span>;
  <span class="pl-smi">int</span> <span class="pl-s1">x_high</span>;

  <span class="pl-k">if</span> (<span class="pl-s1">y_low</span> &gt;= <span class="pl-s1">height</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>) {
    <span class="pl-s1">y_high</span> <span class="pl-c1">=</span> <span class="pl-s1">y_low</span> <span class="pl-c1">=</span> <span class="pl-s1">height</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>;
    <span class="pl-s1">y</span> <span class="pl-c1">=</span> (<span class="pl-smi">T</span>)<span class="pl-s1">y_low</span>;
  } <span class="pl-k">else</span> {
    <span class="pl-s1">y_high</span> <span class="pl-c1">=</span> <span class="pl-s1">y_low</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>;
  }

  <span class="pl-k">if</span> (<span class="pl-s1">x_low</span> &gt;= <span class="pl-s1">width</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>) {
    <span class="pl-s1">x_high</span> <span class="pl-c1">=</span> <span class="pl-s1">x_low</span> <span class="pl-c1">=</span> <span class="pl-s1">width</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>;
    <span class="pl-s1">x</span> <span class="pl-c1">=</span> (<span class="pl-smi">T</span>)<span class="pl-s1">x_low</span>;
  } <span class="pl-k">else</span> {
    <span class="pl-s1">x_high</span> <span class="pl-c1">=</span> <span class="pl-s1">x_low</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>;
  }

  <span class="pl-smi">T</span> <span class="pl-s1">ly</span> <span class="pl-c1">=</span> <span class="pl-s1">y</span> <span class="pl-c1">-</span> <span class="pl-s1">y_low</span>;
  <span class="pl-smi">T</span> <span class="pl-s1">lx</span> <span class="pl-c1">=</span> <span class="pl-s1">x</span> <span class="pl-c1">-</span> <span class="pl-s1">x_low</span>;
  <span class="pl-smi">T</span> <span class="pl-s1">hy</span> <span class="pl-c1">=</span> <span class="pl-c1">1.</span> <span class="pl-c1">-</span> <span class="pl-s1">ly</span>, <span class="pl-s1">hx</span> <span class="pl-c1">=</span> <span class="pl-c1">1.</span> <span class="pl-c1">-</span> <span class="pl-s1">lx</span>;

  <span class="pl-c">// do bilinear interpolation</span>
  <span class="pl-smi">T</span> <span class="pl-s1">v1</span> <span class="pl-c1">=</span> <span class="pl-s1">input</span>[<span class="pl-s1">y_low</span> <span class="pl-c1">*</span> <span class="pl-s1">width</span> <span class="pl-c1">+</span> <span class="pl-s1">x_low</span>];
  <span class="pl-smi">T</span> <span class="pl-s1">v2</span> <span class="pl-c1">=</span> <span class="pl-s1">input</span>[<span class="pl-s1">y_low</span> <span class="pl-c1">*</span> <span class="pl-s1">width</span> <span class="pl-c1">+</span> <span class="pl-s1">x_high</span>];
  <span class="pl-smi">T</span> <span class="pl-s1">v3</span> <span class="pl-c1">=</span> <span class="pl-s1">input</span>[<span class="pl-s1">y_high</span> <span class="pl-c1">*</span> <span class="pl-s1">width</span> <span class="pl-c1">+</span> <span class="pl-s1">x_low</span>];
  <span class="pl-smi">T</span> <span class="pl-s1">v4</span> <span class="pl-c1">=</span> <span class="pl-s1">input</span>[<span class="pl-s1">y_high</span> <span class="pl-c1">*</span> <span class="pl-s1">width</span> <span class="pl-c1">+</span> <span class="pl-s1">x_high</span>];
  <span class="pl-smi">T</span> <span class="pl-s1">w1</span> <span class="pl-c1">=</span> <span class="pl-s1">hy</span> <span class="pl-c1">*</span> <span class="pl-s1">hx</span>, <span class="pl-s1">w2</span> <span class="pl-c1">=</span> <span class="pl-s1">hy</span> <span class="pl-c1">*</span> <span class="pl-s1">lx</span>, <span class="pl-s1">w3</span> <span class="pl-c1">=</span> <span class="pl-s1">ly</span> <span class="pl-c1">*</span> <span class="pl-s1">hx</span>, <span class="pl-s1">w4</span> <span class="pl-c1">=</span> <span class="pl-s1">ly</span> <span class="pl-c1">*</span> <span class="pl-s1">lx</span>;

  <span class="pl-smi">T</span> <span class="pl-s1">val</span> <span class="pl-c1">=</span> (<span class="pl-s1">w1</span> <span class="pl-c1">*</span> <span class="pl-s1">v1</span> <span class="pl-c1">+</span> <span class="pl-s1">w2</span> <span class="pl-c1">*</span> <span class="pl-s1">v2</span> <span class="pl-c1">+</span> <span class="pl-s1">w3</span> <span class="pl-c1">*</span> <span class="pl-s1">v3</span> <span class="pl-c1">+</span> <span class="pl-s1">w4</span> <span class="pl-c1">*</span> <span class="pl-s1">v4</span>);

  <span class="pl-k">return</span> <span class="pl-s1">val</span>;
}

<span class="pl-s1">template</span> <span class="pl-c1">&lt;</span><span class="pl-s1">typename</span> <span class="pl-c1">T</span><span class="pl-c1">&gt;</span>
<span class="pl-s1">void</span> <span class="pl-s1">PSROIAlignForwardCPU</span>(
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">nthreads</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">T</span><span class="pl-c1">*</span> <span class="pl-s1">input</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">T</span> <span class="pl-s1">spatial_scale</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">int</span> <span class="pl-s1">channels</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">int</span> <span class="pl-s1">height</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">int</span> <span class="pl-s1">width</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">int</span> <span class="pl-s1">pooled_height</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">int</span> <span class="pl-s1">pooled_width</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">int</span> <span class="pl-s1">sampling_ratio</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">T</span><span class="pl-c1">*</span> <span class="pl-s1">rois</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">int</span> <span class="pl-s1">channels_out</span>,
    <span class="pl-c1">T</span><span class="pl-c1">*</span> <span class="pl-s1">output</span>,
    <span class="pl-s1">int</span><span class="pl-c1">*</span> <span class="pl-s1">channel_mapping</span>) {
  <span class="pl-smi">int</span> <span class="pl-s1">num_rois</span> <span class="pl-c1">=</span> <span class="pl-s1">nthreads</span> / <span class="pl-s1">channels_out</span> / <span class="pl-s1">pooled_width</span> / <span class="pl-s1">pooled_height</span>;
  <span class="pl-k">for</span> (<span class="pl-smi">int</span> <span class="pl-s1">n</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>; <span class="pl-s1">n</span> <span class="pl-c1">&lt;</span> <span class="pl-s1">num_rois</span>; <span class="pl-s1">n</span><span class="pl-c1">++</span>) {
    <span class="pl-c">// [start, end) interval for spatial sampling</span>
    <span class="pl-k">const</span> <span class="pl-smi">T</span><span class="pl-c1">*</span> <span class="pl-s1">offset_rois</span> <span class="pl-c1">=</span> <span class="pl-s1">rois</span> <span class="pl-c1">+</span> <span class="pl-s1">n</span> <span class="pl-c1">*</span> <span class="pl-c1">5</span>;
    <span class="pl-smi">int</span> <span class="pl-s1">roi_batch_ind</span> <span class="pl-c1">=</span> <span class="pl-s1">offset_rois</span>[<span class="pl-c1">0</span>];

    <span class="pl-c">// Do not using rounding; this implementation detail is critical</span>
    <span class="pl-smi">T</span> <span class="pl-s1">roi_start_w</span> <span class="pl-c1">=</span> <span class="pl-s1">offset_rois</span>[<span class="pl-c1">1</span>] <span class="pl-c1">*</span> <span class="pl-s1">spatial_scale</span> <span class="pl-c1">-</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-c1">0.5</span>);
    <span class="pl-smi">T</span> <span class="pl-s1">roi_start_h</span> <span class="pl-c1">=</span> <span class="pl-s1">offset_rois</span>[<span class="pl-c1">2</span>] <span class="pl-c1">*</span> <span class="pl-s1">spatial_scale</span> <span class="pl-c1">-</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-c1">0.5</span>);
    <span class="pl-smi">T</span> <span class="pl-s1">roi_end_w</span> <span class="pl-c1">=</span> <span class="pl-s1">offset_rois</span>[<span class="pl-c1">3</span>] <span class="pl-c1">*</span> <span class="pl-s1">spatial_scale</span> <span class="pl-c1">-</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-c1">0.5</span>);
    <span class="pl-smi">T</span> <span class="pl-s1">roi_end_h</span> <span class="pl-c1">=</span> <span class="pl-s1">offset_rois</span>[<span class="pl-c1">4</span>] <span class="pl-c1">*</span> <span class="pl-s1">spatial_scale</span> <span class="pl-c1">-</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-c1">0.5</span>);

    <span class="pl-smi">T</span> <span class="pl-s1">roi_width</span> <span class="pl-c1">=</span> <span class="pl-s1">roi_end_w</span> <span class="pl-c1">-</span> <span class="pl-s1">roi_start_w</span>;
    <span class="pl-smi">T</span> <span class="pl-s1">roi_height</span> <span class="pl-c1">=</span> <span class="pl-s1">roi_end_h</span> <span class="pl-c1">-</span> <span class="pl-s1">roi_start_h</span>;
    <span class="pl-smi">T</span> <span class="pl-s1">bin_size_h</span> <span class="pl-c1">=</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">roi_height</span>) / <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">pooled_height</span>);
    <span class="pl-smi">T</span> <span class="pl-s1">bin_size_w</span> <span class="pl-c1">=</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">roi_width</span>) / <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">pooled_width</span>);

    <span class="pl-smi">int</span> <span class="pl-s1">c_in</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>;
    <span class="pl-k">for</span> (<span class="pl-smi">int</span> <span class="pl-s1">c_out</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>; <span class="pl-s1">c_out</span> <span class="pl-c1">&lt;</span> <span class="pl-s1">channels_out</span>; <span class="pl-c1">++</span><span class="pl-s1">c_out</span>) {
      <span class="pl-k">for</span> (<span class="pl-smi">int</span> <span class="pl-s1">ph</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>; <span class="pl-s1">ph</span> <span class="pl-c1">&lt;</span> <span class="pl-s1">pooled_height</span>; <span class="pl-c1">++</span><span class="pl-s1">ph</span>) {
        <span class="pl-k">for</span> (<span class="pl-smi">int</span> <span class="pl-s1">pw</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>; <span class="pl-s1">pw</span> <span class="pl-c1">&lt;</span> <span class="pl-s1">pooled_width</span>; <span class="pl-c1">++</span><span class="pl-s1">pw</span>) {
          <span class="pl-smi">int</span> <span class="pl-s1">index</span> <span class="pl-c1">=</span>
              ((<span class="pl-s1">n</span> <span class="pl-c1">*</span> <span class="pl-s1">channels_out</span> <span class="pl-c1">+</span> <span class="pl-s1">c_out</span>) <span class="pl-c1">*</span> <span class="pl-s1">pooled_height</span> <span class="pl-c1">+</span> <span class="pl-s1">ph</span>) <span class="pl-c1">*</span> <span class="pl-s1">pooled_width</span> <span class="pl-c1">+</span>
              <span class="pl-s1">pw</span>;

          <span class="pl-c">// Do not using floor/ceil; this implementation detail is critical</span>
          <span class="pl-smi">T</span> <span class="pl-s1">hstart</span> <span class="pl-c1">=</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">ph</span>) <span class="pl-c1">*</span> <span class="pl-s1">bin_size_h</span> <span class="pl-c1">+</span> <span class="pl-s1">roi_start_h</span>;
          <span class="pl-smi">T</span> <span class="pl-s1">wstart</span> <span class="pl-c1">=</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">pw</span>) <span class="pl-c1">*</span> <span class="pl-s1">bin_size_w</span> <span class="pl-c1">+</span> <span class="pl-s1">roi_start_w</span>;

          <span class="pl-c">// We use roi_bin_grid to sample the grid and mimic integral</span>
          <span class="pl-smi">int</span> <span class="pl-s1">roi_bin_grid_h</span> <span class="pl-c1">=</span> (<span class="pl-s1">sampling_ratio</span> <span class="pl-c1">&gt;</span> <span class="pl-c1">0</span>)
              ? <span class="pl-s1">sampling_ratio</span>
              : <span class="pl-en">ceil</span>(<span class="pl-s1">roi_height</span> / <span class="pl-s1">pooled_height</span>);
          <span class="pl-smi">int</span> <span class="pl-s1">roi_bin_grid_w</span> <span class="pl-c1">=</span> (<span class="pl-s1">sampling_ratio</span> <span class="pl-c1">&gt;</span> <span class="pl-c1">0</span>)
              ? <span class="pl-s1">sampling_ratio</span>
              : <span class="pl-en">ceil</span>(<span class="pl-s1">roi_width</span> / <span class="pl-s1">pooled_width</span>);
          <span class="pl-k">const</span> <span class="pl-smi">T</span> <span class="pl-s1">count</span> <span class="pl-c1">=</span> <span class="pl-s1">roi_bin_grid_h</span> <span class="pl-c1">*</span> <span class="pl-s1">roi_bin_grid_w</span>;

          <span class="pl-k">const</span> <span class="pl-smi">T</span><span class="pl-c1">*</span> <span class="pl-s1">offset_input</span> <span class="pl-c1">=</span>
              <span class="pl-s1">input</span> <span class="pl-c1">+</span> (<span class="pl-s1">roi_batch_ind</span> <span class="pl-c1">*</span> <span class="pl-s1">channels</span> <span class="pl-c1">+</span> <span class="pl-s1">c_in</span>) <span class="pl-c1">*</span> <span class="pl-s1">height</span> <span class="pl-c1">*</span> <span class="pl-s1">width</span>;

          <span class="pl-smi">T</span> <span class="pl-s1">out_sum</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>;
          <span class="pl-k">for</span> (<span class="pl-smi">int</span> <span class="pl-s1">iy</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>; <span class="pl-s1">iy</span> <span class="pl-c1">&lt;</span> <span class="pl-s1">roi_bin_grid_h</span>; <span class="pl-s1">iy</span><span class="pl-c1">++</span>) {
            <span class="pl-k">const</span> <span class="pl-smi">T</span> <span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">hstart</span> <span class="pl-c1">+</span>
                <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">iy</span> <span class="pl-c1">+</span> <span class="pl-c1">.5f</span>) <span class="pl-c1">*</span> <span class="pl-s1">bin_size_h</span> /
                    <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">roi_bin_grid_h</span>);
            <span class="pl-k">for</span> (<span class="pl-smi">int</span> <span class="pl-s1">ix</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>; <span class="pl-s1">ix</span> <span class="pl-c1">&lt;</span> <span class="pl-s1">roi_bin_grid_w</span>; <span class="pl-s1">ix</span><span class="pl-c1">++</span>) {
              <span class="pl-k">const</span> <span class="pl-smi">T</span> <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">wstart</span> <span class="pl-c1">+</span>
                  <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">ix</span> <span class="pl-c1">+</span> <span class="pl-c1">.5f</span>) <span class="pl-c1">*</span> <span class="pl-s1">bin_size_w</span> /
                      <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">roi_bin_grid_w</span>);
              <span class="pl-smi">T</span> <span class="pl-s1">val</span> <span class="pl-c1">=</span> <span class="pl-en">bilinear_interpolate</span>(
                  <span class="pl-s1">offset_input</span>, <span class="pl-s1">height</span>, <span class="pl-s1">width</span>, <span class="pl-s1">y</span>, <span class="pl-s1">x</span>, <span class="pl-s1">index</span>);
              <span class="pl-s1">out_sum</span> <span class="pl-c1">+=</span> <span class="pl-s1">val</span>;
            }
          }

          <span class="pl-s1">out_sum</span> /= <span class="pl-s1">count</span>;
          <span class="pl-s1">output</span>[<span class="pl-s1">index</span>] <span class="pl-c1">=</span> <span class="pl-s1">out_sum</span>;
          <span class="pl-s1">channel_mapping</span>[<span class="pl-s1">index</span>] <span class="pl-c1">=</span> <span class="pl-s1">c_in</span>;
          <span class="pl-s1">c_in</span><span class="pl-c1">++</span>;
        }
      }
    }
  }
}

<span class="pl-s1">template</span> <span class="pl-c1">&lt;</span><span class="pl-s1">typename</span> <span class="pl-c1">T</span><span class="pl-c1">&gt;</span>
<span class="pl-s1">void</span> <span class="pl-s1">bilinear_interpolate_gradient</span>(
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">height</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">int</span> <span class="pl-s1">width</span>,
    <span class="pl-c1">T</span> <span class="pl-s1">y</span>,
    <span class="pl-c1">T</span> <span class="pl-s1">x</span>,
    <span class="pl-c1">T</span><span class="pl-c1">&amp;</span> <span class="pl-s1">w1</span>,
    <span class="pl-c1">T</span><span class="pl-c1">&amp;</span> <span class="pl-s1">w2</span>,
    <span class="pl-c1">T</span><span class="pl-c1">&amp;</span> <span class="pl-s1">w3</span>,
    <span class="pl-c1">T</span><span class="pl-c1">&amp;</span> <span class="pl-s1">w4</span>,
    <span class="pl-s1">int</span><span class="pl-c1">&amp;</span> <span class="pl-s1">x_low</span>,
    <span class="pl-s1">int</span><span class="pl-c1">&amp;</span> <span class="pl-s1">x_high</span>,
    <span class="pl-s1">int</span><span class="pl-c1">&amp;</span> <span class="pl-s1">y_low</span>,
    <span class="pl-s1">int</span><span class="pl-c1">&amp;</span> <span class="pl-s1">y_high</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">index</span> <span class="pl-c">/* index for debug only*/</span>) {
  <span class="pl-c">// deal with cases that inverse elements are out of feature map boundary</span>
  <span class="pl-k">if</span> (<span class="pl-s1">y</span> <span class="pl-c1">&lt;</span> <span class="pl-c1">-1.0</span> <span class="pl-c1">||</span> <span class="pl-s1">y</span> <span class="pl-c1">&gt;</span> <span class="pl-s1">height</span> <span class="pl-c1">||</span> <span class="pl-s1">x</span> <span class="pl-c1">&lt;</span> <span class="pl-c1">-1.0</span> <span class="pl-c1">||</span> <span class="pl-s1">x</span> <span class="pl-c1">&gt;</span> <span class="pl-s1">width</span>) {
    <span class="pl-c">// empty</span>
    <span class="pl-s1">w1</span> <span class="pl-c1">=</span> <span class="pl-s1">w2</span> <span class="pl-c1">=</span> <span class="pl-s1">w3</span> <span class="pl-c1">=</span> <span class="pl-s1">w4</span> <span class="pl-c1">=</span> <span class="pl-c1">0.</span>;
    <span class="pl-s1">x_low</span> <span class="pl-c1">=</span> <span class="pl-s1">x_high</span> <span class="pl-c1">=</span> <span class="pl-s1">y_low</span> <span class="pl-c1">=</span> <span class="pl-s1">y_high</span> <span class="pl-c1">=</span> <span class="pl-c1">-1</span>;
    <span class="pl-k">return</span>;
  }

  <span class="pl-k">if</span> (<span class="pl-s1">y</span> &lt;= <span class="pl-c1">0</span>)
    <span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>;
  <span class="pl-k">if</span> (<span class="pl-s1">x</span> &lt;= <span class="pl-c1">0</span>)
    <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>;

  <span class="pl-s1">y_low</span> <span class="pl-c1">=</span> (<span class="pl-smi">int</span>)<span class="pl-s1">y</span>;
  <span class="pl-s1">x_low</span> <span class="pl-c1">=</span> (<span class="pl-smi">int</span>)<span class="pl-s1">x</span>;

  <span class="pl-k">if</span> (<span class="pl-s1">y_low</span> &gt;= <span class="pl-s1">height</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>) {
    <span class="pl-s1">y_high</span> <span class="pl-c1">=</span> <span class="pl-s1">y_low</span> <span class="pl-c1">=</span> <span class="pl-s1">height</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>;
    <span class="pl-s1">y</span> <span class="pl-c1">=</span> (<span class="pl-smi">T</span>)<span class="pl-s1">y_low</span>;
  } <span class="pl-k">else</span> {
    <span class="pl-s1">y_high</span> <span class="pl-c1">=</span> <span class="pl-s1">y_low</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>;
  }

  <span class="pl-k">if</span> (<span class="pl-s1">x_low</span> &gt;= <span class="pl-s1">width</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>) {
    <span class="pl-s1">x_high</span> <span class="pl-c1">=</span> <span class="pl-s1">x_low</span> <span class="pl-c1">=</span> <span class="pl-s1">width</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>;
    <span class="pl-s1">x</span> <span class="pl-c1">=</span> (<span class="pl-smi">T</span>)<span class="pl-s1">x_low</span>;
  } <span class="pl-k">else</span> {
    <span class="pl-s1">x_high</span> <span class="pl-c1">=</span> <span class="pl-s1">x_low</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>;
  }

  <span class="pl-smi">T</span> <span class="pl-s1">ly</span> <span class="pl-c1">=</span> <span class="pl-s1">y</span> <span class="pl-c1">-</span> <span class="pl-s1">y_low</span>;
  <span class="pl-smi">T</span> <span class="pl-s1">lx</span> <span class="pl-c1">=</span> <span class="pl-s1">x</span> <span class="pl-c1">-</span> <span class="pl-s1">x_low</span>;
  <span class="pl-smi">T</span> <span class="pl-s1">hy</span> <span class="pl-c1">=</span> <span class="pl-c1">1.</span> <span class="pl-c1">-</span> <span class="pl-s1">ly</span>, <span class="pl-s1">hx</span> <span class="pl-c1">=</span> <span class="pl-c1">1.</span> <span class="pl-c1">-</span> <span class="pl-s1">lx</span>;

  <span class="pl-c">// reference in forward</span>
  <span class="pl-c">// T v1 = input[y_low * width + x_low];</span>
  <span class="pl-c">// T v2 = input[y_low * width + x_high];</span>
  <span class="pl-c">// T v3 = input[y_high * width + x_low];</span>
  <span class="pl-c">// T v4 = input[y_high * width + x_high];</span>
  <span class="pl-c">// T val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);</span>

  <span class="pl-s1">w1</span> <span class="pl-c1">=</span> <span class="pl-s1">hy</span> <span class="pl-c1">*</span> <span class="pl-s1">hx</span>, <span class="pl-s1">w2</span> <span class="pl-c1">=</span> <span class="pl-s1">hy</span> <span class="pl-c1">*</span> <span class="pl-s1">lx</span>, <span class="pl-s1">w3</span> <span class="pl-c1">=</span> <span class="pl-s1">ly</span> <span class="pl-c1">*</span> <span class="pl-s1">hx</span>, <span class="pl-s1">w4</span> <span class="pl-c1">=</span> <span class="pl-s1">ly</span> <span class="pl-c1">*</span> <span class="pl-s1">lx</span>;

  <span class="pl-k">return</span>;
}

<span class="pl-s1">template</span> <span class="pl-c1">&lt;</span><span class="pl-s1">class</span> <span class="pl-c1">T</span><span class="pl-c1">&gt;</span>
<span class="pl-s1">inline</span> <span class="pl-smi">void</span> <span class="pl-en">add</span>(<span class="pl-smi">T</span><span class="pl-c1">*</span> <span class="pl-s1">address</span>, <span class="pl-k">const</span> <span class="pl-smi">T</span><span class="pl-c1">&amp;</span> <span class="pl-s1">val</span>) {
  <span class="pl-c1">*</span><span class="pl-s1">address</span> <span class="pl-c1">+=</span> <span class="pl-s1">val</span>;
}

<span class="pl-s1">template</span> <span class="pl-c1">&lt;</span><span class="pl-s1">typename</span> <span class="pl-c1">T</span><span class="pl-c1">&gt;</span>
<span class="pl-s1">void</span> <span class="pl-s1">PSROIAlignBackwardCPU</span>(
    <span class="pl-k">const</span> <span class="pl-smi">int</span> <span class="pl-s1">nthreads</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">T</span><span class="pl-c1">*</span> <span class="pl-s1">grad_output</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">int</span><span class="pl-c1">*</span> <span class="pl-s1">channel_mapping</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">int</span> <span class="pl-s1">num_rois</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">T</span> <span class="pl-s1">spatial_scale</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">int</span> <span class="pl-s1">channels</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">int</span> <span class="pl-s1">height</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">int</span> <span class="pl-s1">width</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">int</span> <span class="pl-s1">pooled_height</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">int</span> <span class="pl-s1">pooled_width</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">int</span> <span class="pl-s1">sampling_ratio</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">int</span> <span class="pl-s1">channels_out</span>,
    <span class="pl-c1">T</span><span class="pl-c1">*</span> <span class="pl-s1">grad_input</span>,
    <span class="pl-s1">const</span> <span class="pl-smi">T</span><span class="pl-c1">*</span> <span class="pl-s1">rois</span>) {
  <span class="pl-k">for</span> (<span class="pl-smi">int</span> <span class="pl-s1">index</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>; <span class="pl-s1">index</span> <span class="pl-c1">&lt;</span> <span class="pl-s1">nthreads</span>; <span class="pl-s1">index</span><span class="pl-c1">++</span>) {
    <span class="pl-smi">int</span> <span class="pl-s1">pw</span> <span class="pl-c1">=</span> <span class="pl-s1">index</span> % <span class="pl-s1">pooled_width</span>;
    <span class="pl-smi">int</span> <span class="pl-s1">ph</span> <span class="pl-c1">=</span> (<span class="pl-s1">index</span> / <span class="pl-s1">pooled_width</span>) % <span class="pl-s1">pooled_height</span>;
    <span class="pl-smi">int</span> <span class="pl-s1">n</span> <span class="pl-c1">=</span> <span class="pl-s1">index</span> / <span class="pl-s1">pooled_width</span> / <span class="pl-s1">pooled_height</span> / <span class="pl-s1">channels_out</span>;

    <span class="pl-k">const</span> <span class="pl-smi">T</span><span class="pl-c1">*</span> <span class="pl-s1">offset_rois</span> <span class="pl-c1">=</span> <span class="pl-s1">rois</span> <span class="pl-c1">+</span> <span class="pl-s1">n</span> <span class="pl-c1">*</span> <span class="pl-c1">5</span>;
    <span class="pl-smi">int</span> <span class="pl-s1">roi_batch_ind</span> <span class="pl-c1">=</span> <span class="pl-s1">offset_rois</span>[<span class="pl-c1">0</span>];

    <span class="pl-c">// Do not using rounding; this implementation detail is critical</span>
    <span class="pl-smi">T</span> <span class="pl-s1">roi_start_w</span> <span class="pl-c1">=</span> <span class="pl-s1">offset_rois</span>[<span class="pl-c1">1</span>] <span class="pl-c1">*</span> <span class="pl-s1">spatial_scale</span> <span class="pl-c1">-</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-c1">0.5</span>);
    <span class="pl-smi">T</span> <span class="pl-s1">roi_start_h</span> <span class="pl-c1">=</span> <span class="pl-s1">offset_rois</span>[<span class="pl-c1">2</span>] <span class="pl-c1">*</span> <span class="pl-s1">spatial_scale</span> <span class="pl-c1">-</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-c1">0.5</span>);
    <span class="pl-smi">T</span> <span class="pl-s1">roi_end_w</span> <span class="pl-c1">=</span> <span class="pl-s1">offset_rois</span>[<span class="pl-c1">3</span>] <span class="pl-c1">*</span> <span class="pl-s1">spatial_scale</span> <span class="pl-c1">-</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-c1">0.5</span>);
    <span class="pl-smi">T</span> <span class="pl-s1">roi_end_h</span> <span class="pl-c1">=</span> <span class="pl-s1">offset_rois</span>[<span class="pl-c1">4</span>] <span class="pl-c1">*</span> <span class="pl-s1">spatial_scale</span> <span class="pl-c1">-</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-c1">0.5</span>);

    <span class="pl-c">// Force too small ROIs to be 1x1</span>
    <span class="pl-smi">T</span> <span class="pl-s1">roi_width</span> <span class="pl-c1">=</span> <span class="pl-s1">roi_end_w</span> <span class="pl-c1">-</span> <span class="pl-s1">roi_start_w</span>;
    <span class="pl-smi">T</span> <span class="pl-s1">roi_height</span> <span class="pl-c1">=</span> <span class="pl-s1">roi_end_h</span> <span class="pl-c1">-</span> <span class="pl-s1">roi_start_h</span>;
    <span class="pl-smi">T</span> <span class="pl-s1">bin_size_h</span> <span class="pl-c1">=</span> <span class="pl-s1">roi_height</span> / <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">pooled_height</span>);
    <span class="pl-smi">T</span> <span class="pl-s1">bin_size_w</span> <span class="pl-c1">=</span> <span class="pl-s1">roi_width</span> / <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">pooled_width</span>);

    <span class="pl-smi">int</span> <span class="pl-s1">c_in</span> <span class="pl-c1">=</span> <span class="pl-s1">channel_mapping</span>[<span class="pl-s1">index</span>];
    <span class="pl-smi">T</span><span class="pl-c1">*</span> <span class="pl-s1">grad_input_offset</span> <span class="pl-c1">=</span>
        <span class="pl-s1">grad_input</span> <span class="pl-c1">+</span> (<span class="pl-s1">roi_batch_ind</span> <span class="pl-c1">*</span> <span class="pl-s1">channels</span> <span class="pl-c1">+</span> <span class="pl-s1">c_in</span>) <span class="pl-c1">*</span> <span class="pl-s1">height</span> <span class="pl-c1">*</span> <span class="pl-s1">width</span>;

    <span class="pl-c">// Do not using floor/ceil; this implementation detail is critical</span>
    <span class="pl-smi">T</span> <span class="pl-s1">hstart</span> <span class="pl-c1">=</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">ph</span>) <span class="pl-c1">*</span> <span class="pl-s1">bin_size_h</span> <span class="pl-c1">+</span> <span class="pl-s1">roi_start_h</span>;
    <span class="pl-smi">T</span> <span class="pl-s1">wstart</span> <span class="pl-c1">=</span> <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">pw</span>) <span class="pl-c1">*</span> <span class="pl-s1">bin_size_w</span> <span class="pl-c1">+</span> <span class="pl-s1">roi_start_w</span>;

    <span class="pl-k">const</span> <span class="pl-smi">T</span> <span class="pl-s1">grad_output_this_bin</span> <span class="pl-c1">=</span> <span class="pl-s1">grad_output</span>[<span class="pl-s1">index</span>];

    <span class="pl-c">// We use roi_bin_grid to sample the grid and mimic integral</span>
    <span class="pl-smi">int</span> <span class="pl-s1">roi_bin_grid_h</span> <span class="pl-c1">=</span> (<span class="pl-s1">sampling_ratio</span> <span class="pl-c1">&gt;</span> <span class="pl-c1">0</span>)
        ? <span class="pl-s1">sampling_ratio</span>
        : <span class="pl-en">ceil</span>(<span class="pl-s1">roi_height</span> / <span class="pl-s1">pooled_height</span>); <span class="pl-c">// e.g., = 2</span>
    <span class="pl-smi">int</span> <span class="pl-s1">roi_bin_grid_w</span> <span class="pl-c1">=</span>
        (<span class="pl-s1">sampling_ratio</span> <span class="pl-c1">&gt;</span> <span class="pl-c1">0</span>) ? <span class="pl-s1">sampling_ratio</span> : <span class="pl-en">ceil</span>(<span class="pl-s1">roi_width</span> / <span class="pl-s1">pooled_width</span>);
    <span class="pl-k">const</span> <span class="pl-smi">T</span> <span class="pl-s1">count</span> <span class="pl-c1">=</span> <span class="pl-s1">roi_bin_grid_h</span> <span class="pl-c1">*</span> <span class="pl-s1">roi_bin_grid_w</span>;

    <span class="pl-k">for</span> (<span class="pl-smi">int</span> <span class="pl-s1">iy</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>; <span class="pl-s1">iy</span> <span class="pl-c1">&lt;</span> <span class="pl-s1">roi_bin_grid_h</span>; <span class="pl-s1">iy</span><span class="pl-c1">++</span>) {
      <span class="pl-k">const</span> <span class="pl-smi">T</span> <span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">hstart</span> <span class="pl-c1">+</span>
          <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">iy</span> <span class="pl-c1">+</span> <span class="pl-c1">.5f</span>) <span class="pl-c1">*</span> <span class="pl-s1">bin_size_h</span> /
              <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">roi_bin_grid_h</span>);
      <span class="pl-k">for</span> (<span class="pl-smi">int</span> <span class="pl-s1">ix</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>; <span class="pl-s1">ix</span> <span class="pl-c1">&lt;</span> <span class="pl-s1">roi_bin_grid_w</span>; <span class="pl-s1">ix</span><span class="pl-c1">++</span>) {
        <span class="pl-k">const</span> <span class="pl-smi">T</span> <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">wstart</span> <span class="pl-c1">+</span>
            <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">ix</span> <span class="pl-c1">+</span> <span class="pl-c1">.5f</span>) <span class="pl-c1">*</span> <span class="pl-s1">bin_size_w</span> /
                <span class="pl-s1">static_cast</span><span class="pl-c1">&lt;</span><span class="pl-c1">T</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">roi_bin_grid_w</span>);

        <span class="pl-smi">T</span> <span class="pl-s1">w1</span>, <span class="pl-s1">w2</span>, <span class="pl-s1">w3</span>, <span class="pl-s1">w4</span>;
        <span class="pl-smi">int</span> <span class="pl-s1">x_low</span>, <span class="pl-s1">x_high</span>, <span class="pl-s1">y_low</span>, <span class="pl-s1">y_high</span>;

        <span class="pl-en">bilinear_interpolate_gradient</span>(
            <span class="pl-s1">height</span>,
            <span class="pl-s1">width</span>,
            <span class="pl-s1">y</span>,
            <span class="pl-s1">x</span>,
            <span class="pl-s1">w1</span>,
            <span class="pl-s1">w2</span>,
            <span class="pl-s1">w3</span>,
            <span class="pl-s1">w4</span>,
            <span class="pl-s1">x_low</span>,
            <span class="pl-s1">x_high</span>,
            <span class="pl-s1">y_low</span>,
            <span class="pl-s1">y_high</span>,
            <span class="pl-s1">index</span>);

        <span class="pl-smi">T</span> <span class="pl-s1">g1</span> <span class="pl-c1">=</span> <span class="pl-s1">grad_output_this_bin</span> <span class="pl-c1">*</span> <span class="pl-s1">w1</span> / <span class="pl-s1">count</span>;
        <span class="pl-smi">T</span> <span class="pl-s1">g2</span> <span class="pl-c1">=</span> <span class="pl-s1">grad_output_this_bin</span> <span class="pl-c1">*</span> <span class="pl-s1">w2</span> / <span class="pl-s1">count</span>;
        <span class="pl-smi">T</span> <span class="pl-s1">g3</span> <span class="pl-c1">=</span> <span class="pl-s1">grad_output_this_bin</span> <span class="pl-c1">*</span> <span class="pl-s1">w3</span> / <span class="pl-s1">count</span>;
        <span class="pl-smi">T</span> <span class="pl-s1">g4</span> <span class="pl-c1">=</span> <span class="pl-s1">grad_output_this_bin</span> <span class="pl-c1">*</span> <span class="pl-s1">w4</span> / <span class="pl-s1">count</span>;

        <span class="pl-k">if</span> (<span class="pl-s1">x_low</span> &gt;= <span class="pl-c1">0</span> <span class="pl-c1">&amp;&amp;</span> <span class="pl-s1">x_high</span> &gt;= <span class="pl-c1">0</span> <span class="pl-c1">&amp;&amp;</span> <span class="pl-s1">y_low</span> &gt;= <span class="pl-c1">0</span> <span class="pl-c1">&amp;&amp;</span> <span class="pl-s1">y_high</span> &gt;= <span class="pl-c1">0</span>) {
          <span class="pl-en">add</span>(<span class="pl-s1">grad_input_offset</span> <span class="pl-c1">+</span> <span class="pl-s1">y_low</span> <span class="pl-c1">*</span> <span class="pl-s1">width</span> <span class="pl-c1">+</span> <span class="pl-s1">x_low</span>, <span class="pl-s1">g1</span>);
          <span class="pl-en">add</span>(<span class="pl-s1">grad_input_offset</span> <span class="pl-c1">+</span> <span class="pl-s1">y_low</span> <span class="pl-c1">*</span> <span class="pl-s1">width</span> <span class="pl-c1">+</span> <span class="pl-s1">x_high</span>, <span class="pl-s1">g2</span>);
          <span class="pl-en">add</span>(<span class="pl-s1">grad_input_offset</span> <span class="pl-c1">+</span> <span class="pl-s1">y_high</span> <span class="pl-c1">*</span> <span class="pl-s1">width</span> <span class="pl-c1">+</span> <span class="pl-s1">x_low</span>, <span class="pl-s1">g3</span>);
          <span class="pl-en">add</span>(<span class="pl-s1">grad_input_offset</span> <span class="pl-c1">+</span> <span class="pl-s1">y_high</span> <span class="pl-c1">*</span> <span class="pl-s1">width</span> <span class="pl-c1">+</span> <span class="pl-s1">x_high</span>, <span class="pl-s1">g4</span>);
        } <span class="pl-c">// if</span>
      } <span class="pl-c">// ix</span>
    } <span class="pl-c">// iy</span>
  }
}

<span class="pl-ent">std</span>::<span class="pl-s1">tuple</span><span class="pl-c1">&lt;</span><span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span>, <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&gt;</span> <span class="pl-en">PSROIAlign_forward_cpu</span>(
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">input</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">rois</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">float</span> <span class="pl-s1">spatial_scale</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">pooled_height</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">pooled_width</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">sampling_ratio</span>) {
  <span class="pl-c">// Check if input tensors are CPU tensors</span>
  <span class="pl-en">AT_ASSERTM</span>(<span class="pl-s1">input</span>.<span class="pl-en">device</span>().<span class="pl-en">is_cpu</span>(), <span class="pl-s">"input must be a CPU tensor"</span>);
  <span class="pl-en">AT_ASSERTM</span>(<span class="pl-s1">rois</span>.<span class="pl-en">device</span>().<span class="pl-en">is_cpu</span>(), <span class="pl-s">"rois must be a CPU tensor"</span>);

  <span class="pl-s1">at</span>::<span class="pl-s1">TensorArg</span> <span class="pl-s1">input_t</span>{<span class="pl-s1">input</span>, <span class="pl-s">"input"</span>, <span class="pl-c1">1</span>}, <span class="pl-s1">rois_t</span>{<span class="pl-s1">rois</span>, <span class="pl-s">"rois"</span>, <span class="pl-c1">2</span>};

  <span class="pl-ent">at</span>::<span class="pl-s1">CheckedFrom</span> <span class="pl-s1">c</span> <span class="pl-c1">=</span> <span class="pl-s">"PSROIAlign_forward_cpu"</span>;
  <span class="pl-ent">at</span>::<span class="pl-s1">checkAllSameType</span>(<span class="pl-s1">c</span>, {<span class="pl-s1">input_t</span>, <span class="pl-s1">rois_t</span>});

  <span class="pl-smi">int</span> <span class="pl-s1">num_rois</span> <span class="pl-c1">=</span> <span class="pl-s1">rois</span>.<span class="pl-en">size</span>(<span class="pl-c1">0</span>);
  <span class="pl-smi">int</span> <span class="pl-s1">channels</span> <span class="pl-c1">=</span> <span class="pl-s1">input</span>.<span class="pl-en">size</span>(<span class="pl-c1">1</span>);
  <span class="pl-smi">int</span> <span class="pl-s1">height</span> <span class="pl-c1">=</span> <span class="pl-s1">input</span>.<span class="pl-en">size</span>(<span class="pl-c1">2</span>);
  <span class="pl-smi">int</span> <span class="pl-s1">width</span> <span class="pl-c1">=</span> <span class="pl-s1">input</span>.<span class="pl-en">size</span>(<span class="pl-c1">3</span>);

  <span class="pl-en">AT_ASSERTM</span>(
      <span class="pl-s1">channels</span> % (<span class="pl-s1">pooled_height</span> <span class="pl-c1">*</span> <span class="pl-s1">pooled_width</span>) <span class="pl-c1">==</span> <span class="pl-c1">0</span>,
      <span class="pl-s">"input channels must be a multiple of pooling height * pooling width"</span>);
  <span class="pl-smi">int</span> <span class="pl-s1">channels_out</span> <span class="pl-c1">=</span> <span class="pl-s1">channels</span> / (<span class="pl-s1">pooled_height</span> <span class="pl-c1">*</span> <span class="pl-s1">pooled_width</span>);

  auto <span class="pl-smi">output</span> <span class="pl-c1">=</span> <span class="pl-s1">at</span>::<span class="pl-en">zeros</span>(
      {<span class="pl-smi">num_rois</span>, <span class="pl-smi">channels_out</span>, <span class="pl-smi">pooled_height</span>, <span class="pl-smi">pooled_width</span>}, <span class="pl-smi">input</span>.<span class="pl-en">options</span>());
  auto <span class="pl-smi">channel_mapping</span> <span class="pl-c1">=</span>
      <span class="pl-s1">at</span>::<span class="pl-en">zeros</span>(<span class="pl-smi">output</span>.<span class="pl-en">sizes</span>(), <span class="pl-smi">input</span>.<span class="pl-en">options</span>().<span class="pl-en">dtype</span>(<span class="pl-smi">at</span>::<span class="pl-s1">kInt</span>));

  auto <span class="pl-smi">output_size</span><span class="pl-s1"></span> <span class="pl-c1">=</span> <span class="pl-s1">output</span>.<span class="pl-en">numel</span>();
  <span class="pl-k">if</span> (<span class="pl-s1">output_size</span> <span class="pl-c1">==</span> <span class="pl-c1">0</span>) {
    <span class="pl-k">return</span> <span class="pl-s1">std</span>::<span class="pl-en">make_tuple</span>(<span class="pl-s1">output</span>, <span class="pl-s1">channel_mapping</span>);
  }

  <span class="pl-en">AT_DISPATCH_FLOATING_TYPES_AND_HALF</span>(
      <span class="pl-s1">input</span>.<span class="pl-en">scalar_type</span>(), <span class="pl-s">"PSROIAlign_forward"</span>, [<span class="pl-c1">&amp;</span>] {
        <span class="pl-s1">PSROIAlignForwardCPU</span><span class="pl-c1">&lt;</span><span class="pl-s1">scalar_t</span><span class="pl-c1">&gt;</span>(
            <span class="pl-s1">output_size</span>,
            <span class="pl-s1">input</span>.<span class="pl-en">contiguous</span>().<span class="pl-c1">data_ptr</span><span class="pl-c1">&lt;</span><span class="pl-s1">scalar_t</span><span class="pl-c1">&gt;</span>(<span class="pl-s1"></span>),
            <span class="pl-s1">spatial_scale</span>,
            <span class="pl-s1">channels</span>,
            <span class="pl-s1">height</span>,
            <span class="pl-s1">width</span>,
            <span class="pl-s1">pooled_height</span>,
            <span class="pl-s1">pooled_width</span>,
            <span class="pl-s1">sampling_ratio</span>,
            <span class="pl-s1">rois</span>.<span class="pl-en">contiguous</span>().<span class="pl-c1">data_ptr</span><span class="pl-c1">&lt;</span><span class="pl-s1">scalar_t</span><span class="pl-c1">&gt;</span>(<span class="pl-s1"></span>),
            <span class="pl-s1">channels_out</span>,
            <span class="pl-s1">output</span>.<span class="pl-c1">data_ptr</span><span class="pl-c1">&lt;</span><span class="pl-s1">scalar_t</span><span class="pl-c1">&gt;</span>(<span class="pl-s1"></span>),
            <span class="pl-s1">channel_mapping</span>.<span class="pl-c1">data_ptr</span><span class="pl-c1">&lt;</span><span class="pl-s1">int</span><span class="pl-c1">&gt;</span>(<span class="pl-s1"></span>));
      });
  <span class="pl-k">return</span> <span class="pl-s1">std</span>::<span class="pl-en">make_tuple</span>(<span class="pl-s1">output</span>, <span class="pl-s1">channel_mapping</span>);
}

<span class="pl-ent">at</span>::<span class="pl-s1">Tensor</span> <span class="pl-en">PSROIAlign_backward_cpu</span>(
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">grad</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">rois</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">channel_mapping</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">float</span> <span class="pl-s1">spatial_scale</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">pooled_height</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">pooled_width</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">sampling_ratio</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">batch_size</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">channels</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">height</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">width</span>) {
  <span class="pl-c">// Check if input tensors are CPU tensors</span>
  <span class="pl-en">AT_ASSERTM</span>(<span class="pl-s1">grad</span>.<span class="pl-en">device</span>().<span class="pl-en">is_cpu</span>(), <span class="pl-s">"grad must be a CPU tensor"</span>);
  <span class="pl-en">AT_ASSERTM</span>(<span class="pl-s1">rois</span>.<span class="pl-en">device</span>().<span class="pl-en">is_cpu</span>(), <span class="pl-s">"rois must be a CPU tensor"</span>);
  <span class="pl-en">AT_ASSERTM</span>(
      <span class="pl-s1">channel_mapping</span>.<span class="pl-en">device</span>().<span class="pl-en">is_cpu</span>(),
      <span class="pl-s">"channel_mapping must be a CPU tensor"</span>);

  <span class="pl-s1">at</span>::<span class="pl-s1">TensorArg</span> <span class="pl-s1">grad_t</span>{<span class="pl-s1">grad</span>, <span class="pl-s">"grad"</span>, <span class="pl-c1">1</span>}, <span class="pl-s1">rois_t</span>{<span class="pl-s1">rois</span>, <span class="pl-s">"rois"</span>, <span class="pl-c1">2</span>},
      <span class="pl-s1">channel_mapping_t</span>{<span class="pl-s1">channel_mapping</span>, <span class="pl-s">"channel_mapping"</span>, <span class="pl-c1">3</span>};

  <span class="pl-ent">at</span>::<span class="pl-s1">CheckedFrom</span> <span class="pl-s1">c</span> <span class="pl-c1">=</span> <span class="pl-s">"PSROIAlign_backward_cpu"</span>;
  <span class="pl-ent">at</span>::<span class="pl-s1">checkAllSameType</span>(<span class="pl-s1">c</span>, {<span class="pl-s1">grad_t</span>, <span class="pl-s1">rois_t</span>});

  auto <span class="pl-smi">num_rois</span><span class="pl-s1"></span> <span class="pl-c1">=</span> <span class="pl-s1">rois</span>.<span class="pl-en">size</span>(<span class="pl-c1">0</span>);
  auto <span class="pl-smi">grad_input</span> <span class="pl-c1">=</span>
      <span class="pl-s1">at</span>::<span class="pl-en">zeros</span>({<span class="pl-smi">batch_size</span>, <span class="pl-smi">channels</span>, <span class="pl-smi">height</span>, <span class="pl-smi">width</span>}, <span class="pl-smi">grad</span>.<span class="pl-en">options</span>());

  <span class="pl-c">// handle possibly empty gradients</span>
  <span class="pl-k">if</span> (<span class="pl-s1">grad</span>.<span class="pl-en">numel</span>() <span class="pl-c1">==</span> <span class="pl-c1">0</span>) {
    <span class="pl-k">return</span> <span class="pl-s1">grad_input</span>;
  }

  <span class="pl-smi">int</span> <span class="pl-s1">channels_out</span> <span class="pl-c1">=</span> <span class="pl-s1">channels</span> / (<span class="pl-s1">pooled_height</span> <span class="pl-c1">*</span> <span class="pl-s1">pooled_width</span>);

  <span class="pl-en">AT_DISPATCH_FLOATING_TYPES_AND_HALF</span>(
      <span class="pl-s1">grad</span>.<span class="pl-en">scalar_type</span>(), <span class="pl-s">"PSROIAlign_backward"</span>, [<span class="pl-c1">&amp;</span>] {
        <span class="pl-s1">PSROIAlignBackwardCPU</span><span class="pl-c1">&lt;</span><span class="pl-s1">scalar_t</span><span class="pl-c1">&gt;</span>(
            <span class="pl-s1">grad</span>.<span class="pl-en">numel</span>(),
            <span class="pl-s1">grad</span>.<span class="pl-en">contiguous</span>().<span class="pl-c1">data_ptr</span><span class="pl-c1">&lt;</span><span class="pl-s1">scalar_t</span><span class="pl-c1">&gt;</span>(<span class="pl-s1"></span>),
            <span class="pl-s1">channel_mapping</span>.<span class="pl-c1">data_ptr</span><span class="pl-c1">&lt;</span><span class="pl-s1">int</span><span class="pl-c1">&gt;</span>(<span class="pl-s1"></span>),
            <span class="pl-s1">num_rois</span>,
            <span class="pl-s1">spatial_scale</span>,
            <span class="pl-s1">channels</span>,
            <span class="pl-s1">height</span>,
            <span class="pl-s1">width</span>,
            <span class="pl-s1">pooled_height</span>,
            <span class="pl-s1">pooled_width</span>,
            <span class="pl-s1">sampling_ratio</span>,
            <span class="pl-s1">channels_out</span>,
            <span class="pl-s1">grad_input</span>.<span class="pl-c1">data_ptr</span><span class="pl-c1">&lt;</span><span class="pl-s1">scalar_t</span><span class="pl-c1">&gt;</span>(<span class="pl-s1"></span>),
            <span class="pl-s1">rois</span>.<span class="pl-en">contiguous</span>().<span class="pl-c1">data_ptr</span><span class="pl-c1">&lt;</span><span class="pl-s1">scalar_t</span><span class="pl-c1">&gt;</span>(<span class="pl-s1"></span>));
      });
  <span class="pl-k">return</span> <span class="pl-s1">grad_input</span>;
}</pre></div>
</details>
<details>
<summary>
PSROIAlign.h
</summary>
<div class="highlight highlight-source-c notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="#pragma once

#include &quot;torch/torch.h&quot;
#include &quot;PSROIAlign_cuda.h&quot;

#include &lt;iostream&gt;



std::tuple&lt;at::Tensor, at::Tensor&gt; PSROIAlign_forward_cpu(
    const at::Tensor &amp; input,
    const at::Tensor &amp; rois,
    const float spatial_scale,
    const int pooled_height,
    const int pooled_width,
    const int sampling_ratio);

at::Tensor PSROIAlign_backward_cpu(
    const at::Tensor &amp; grad,
    const at::Tensor &amp; rois,
    const at::Tensor &amp; mapping_channel,
    const float spatial_scale,
    const int pooled_height,
    const int pooled_width,
    const int sampling_ratio,
    const int batch_size,
    const int channels,
    const int height,
    const int width);

std::tuple&lt;at::Tensor, at::Tensor&gt; PSROIAlign_forward(
    const at::Tensor&amp; input,
    const at::Tensor&amp; rois,
    const float spatial_scale,
    const int pooled_height,
    const int pooled_width,
    const int sampling_ratio) {
  if (input.type().is_cuda()) {

    return PSROIAlign_forward_cuda(
        input,
        rois,
        spatial_scale,
        pooled_height,
        pooled_width,
        sampling_ratio);
  }
  return PSROIAlign_forward_cpu(
      input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio);
}

at::Tensor PSROIAlign_backward(
    const at::Tensor&amp; grad,
    const at::Tensor&amp; rois,
    const at::Tensor&amp; mapping_channel,
    const float spatial_scale,
    const int pooled_height,
    const int pooled_width,
    const int sampling_ratio,
    const int batch_size,
    const int channels,
    const int height,
    const int width) {
  if (grad.type().is_cuda()) {

    return PSROIAlign_backward_cuda(
        grad,
        rois,
        mapping_channel,
        spatial_scale,
        pooled_height,
        pooled_width,
        sampling_ratio,
        batch_size,
        channels,
        height,
        width);
  }
  return PSROIAlign_backward_cpu(
      grad,
      rois,
      mapping_channel,
      spatial_scale,
      pooled_height,
      pooled_width,
      sampling_ratio,
      batch_size,
      channels,
      height,
      width);
}

using namespace at;
using torch::Tensor;
using torch::autograd::AutogradContext;
using torch::autograd::Variable;
using torch::autograd::variable_list;

class PSROIAlignFunction
    : public torch::autograd::Function&lt;PSROIAlignFunction&gt; {
 public:
  static variable_list forward(
      AutogradContext* ctx,
      Variable input,
      Variable rois,
      const double spatial_scale,
      const int64_t pooled_height,
      const int64_t pooled_width,
      const int64_t sampling_ratio) {
    ctx-&gt;saved_data[&quot;spatial_scale&quot;] = spatial_scale;
    ctx-&gt;saved_data[&quot;pooled_height&quot;] = pooled_height;
    ctx-&gt;saved_data[&quot;pooled_width&quot;] = pooled_width;
    ctx-&gt;saved_data[&quot;sampling_ratio&quot;] = sampling_ratio;
    ctx-&gt;saved_data[&quot;input_shape&quot;] = input.sizes();
    auto result = PSROIAlign_forward(
        input,
        rois,
        spatial_scale,
        pooled_height,
        pooled_width,
        sampling_ratio);
    auto output = std::get&lt;0&gt;(result);
    auto channel_mapping = std::get&lt;1&gt;(result);
    ctx-&gt;save_for_backward({rois, channel_mapping});
    ctx-&gt;mark_non_differentiable({channel_mapping});
    return {output, channel_mapping};
  }

  static variable_list backward(
      AutogradContext* ctx,
      variable_list grad_output) {
    // Use data saved in forward
    auto saved = ctx-&gt;get_saved_variables();
    auto rois = saved[0];
    auto channel_mapping = saved[1];
    auto input_shape = ctx-&gt;saved_data[&quot;input_shape&quot;].toIntList();
    auto grad_in = PSROIAlign_backward(
        grad_output[0],
        rois,
        channel_mapping,
        ctx-&gt;saved_data[&quot;spatial_scale&quot;].toDouble(),
        ctx-&gt;saved_data[&quot;pooled_height&quot;].toInt(),
        ctx-&gt;saved_data[&quot;pooled_width&quot;].toInt(),
        ctx-&gt;saved_data[&quot;sampling_ratio&quot;].toInt(),
        input_shape[0],
        input_shape[1],
        input_shape[2],
        input_shape[3]);
    return {
        grad_in, Variable(), Variable(), Variable(), Variable(), Variable()};
  }
};

std::tuple&lt;Tensor, Tensor&gt; ps_roi_align(
    const Tensor&amp; input,
    const torch::Tensor&amp; rois,
    const double spatial_scale,
    const int64_t pooled_height,
    const int64_t pooled_width,
    const int64_t sampling_ratio) {
  auto result = PSROIAlignFunction::apply(
      input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio);
  return std::tuple&lt;Tensor, Tensor&gt;(result[0], result[1]);
}"><pre class="notranslate"><span class="pl-k">#pragma</span> once

<span class="pl-k">#include</span> <span class="pl-s">"torch/torch.h"</span>
<span class="pl-k">#include</span> <span class="pl-s">"PSROIAlign_cuda.h"</span>

<span class="pl-k">#include</span> <span class="pl-s">&lt;iostream&gt;</span>



<span class="pl-ent">std</span>::<span class="pl-s1">tuple</span><span class="pl-c1">&lt;</span><span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span>, <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&gt;</span> <span class="pl-en">PSROIAlign_forward_cpu</span>(
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span> <span class="pl-c1">&amp;</span> <span class="pl-s1">input</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span> <span class="pl-c1">&amp;</span> <span class="pl-s1">rois</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">float</span> <span class="pl-s1">spatial_scale</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">pooled_height</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">pooled_width</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">sampling_ratio</span>);

<span class="pl-ent">at</span>::<span class="pl-s1">Tensor</span> <span class="pl-en">PSROIAlign_backward_cpu</span>(
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span> <span class="pl-c1">&amp;</span> <span class="pl-s1">grad</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span> <span class="pl-c1">&amp;</span> <span class="pl-s1">rois</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span> <span class="pl-c1">&amp;</span> <span class="pl-s1">mapping_channel</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">float</span> <span class="pl-s1">spatial_scale</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">pooled_height</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">pooled_width</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">sampling_ratio</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">batch_size</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">channels</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">height</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">width</span>);

<span class="pl-ent">std</span>::<span class="pl-s1">tuple</span><span class="pl-c1">&lt;</span><span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span>, <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&gt;</span> <span class="pl-en">PSROIAlign_forward</span>(
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">input</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">rois</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">float</span> <span class="pl-s1">spatial_scale</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">pooled_height</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">pooled_width</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">sampling_ratio</span>) {
  <span class="pl-k">if</span> (<span class="pl-s1">input</span>.<span class="pl-en">type</span>().<span class="pl-en">is_cuda</span>()) {

    <span class="pl-k">return</span> <span class="pl-en">PSROIAlign_forward_cuda</span>(
        <span class="pl-s1">input</span>,
        <span class="pl-s1">rois</span>,
        <span class="pl-s1">spatial_scale</span>,
        <span class="pl-s1">pooled_height</span>,
        <span class="pl-s1">pooled_width</span>,
        <span class="pl-s1">sampling_ratio</span>);
  }
  <span class="pl-k">return</span> <span class="pl-en">PSROIAlign_forward_cpu</span>(
      <span class="pl-s1">input</span>, <span class="pl-s1">rois</span>, <span class="pl-s1">spatial_scale</span>, <span class="pl-s1">pooled_height</span>, <span class="pl-s1">pooled_width</span>, <span class="pl-s1">sampling_ratio</span>);
}

<span class="pl-ent">at</span>::<span class="pl-s1">Tensor</span> <span class="pl-en">PSROIAlign_backward</span>(
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">grad</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">rois</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">at</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">mapping_channel</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">float</span> <span class="pl-s1">spatial_scale</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">pooled_height</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">pooled_width</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">sampling_ratio</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">batch_size</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">channels</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">height</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int</span> <span class="pl-s1">width</span>) {
  <span class="pl-k">if</span> (<span class="pl-s1">grad</span>.<span class="pl-en">type</span>().<span class="pl-en">is_cuda</span>()) {

    <span class="pl-k">return</span> <span class="pl-en">PSROIAlign_backward_cuda</span>(
        <span class="pl-s1">grad</span>,
        <span class="pl-s1">rois</span>,
        <span class="pl-s1">mapping_channel</span>,
        <span class="pl-s1">spatial_scale</span>,
        <span class="pl-s1">pooled_height</span>,
        <span class="pl-s1">pooled_width</span>,
        <span class="pl-s1">sampling_ratio</span>,
        <span class="pl-s1">batch_size</span>,
        <span class="pl-s1">channels</span>,
        <span class="pl-s1">height</span>,
        <span class="pl-s1">width</span>);
  }
  <span class="pl-k">return</span> <span class="pl-en">PSROIAlign_backward_cpu</span>(
      <span class="pl-s1">grad</span>,
      <span class="pl-s1">rois</span>,
      <span class="pl-s1">mapping_channel</span>,
      <span class="pl-s1">spatial_scale</span>,
      <span class="pl-s1">pooled_height</span>,
      <span class="pl-s1">pooled_width</span>,
      <span class="pl-s1">sampling_ratio</span>,
      <span class="pl-s1">batch_size</span>,
      <span class="pl-s1">channels</span>,
      <span class="pl-s1">height</span>,
      <span class="pl-s1">width</span>);
}

<span class="pl-smi">using</span> <span class="pl-s1">namespace</span> <span class="pl-s1">at</span>;
<span class="pl-smi">using</span> <span class="pl-s1">torch</span>::<span class="pl-s1">Tensor</span>;
<span class="pl-smi">using</span> <span class="pl-s1">torch</span>::<span class="pl-s1">autograd</span>::<span class="pl-s1">AutogradContext</span>;
<span class="pl-smi">using</span> <span class="pl-s1">torch</span>::<span class="pl-s1">autograd</span>::<span class="pl-s1">Variable</span>;
<span class="pl-smi">using</span> <span class="pl-s1">torch</span>::<span class="pl-s1">autograd</span>::<span class="pl-s1">variable_list</span>;

<span class="pl-smi">class</span> <span class="pl-s1">PSROIAlignFunction</span>
    : <span class="pl-s1">public</span> <span class="pl-s1">torch</span>::<span class="pl-s1">autograd</span>::<span class="pl-s1">Function</span><span class="pl-c1">&lt;</span><span class="pl-s1">PSROIAlignFunction</span><span class="pl-c1">&gt;</span> {
 <span class="pl-ent">public</span>:
  <span class="pl-s1">static</span> <span class="pl-smi">variable_list</span> <span class="pl-en">forward</span>(
      <span class="pl-smi">AutogradContext</span><span class="pl-c1">*</span> <span class="pl-s1">ctx</span>,
      <span class="pl-smi">Variable</span> <span class="pl-s1">input</span>,
      <span class="pl-smi">Variable</span> <span class="pl-s1">rois</span>,
      <span class="pl-k">const</span> <span class="pl-smi">double</span> <span class="pl-s1">spatial_scale</span>,
      <span class="pl-k">const</span> <span class="pl-smi">int64_t</span> <span class="pl-s1">pooled_height</span>,
      <span class="pl-k">const</span> <span class="pl-smi">int64_t</span> <span class="pl-s1">pooled_width</span>,
      <span class="pl-k">const</span> <span class="pl-smi">int64_t</span> <span class="pl-s1">sampling_ratio</span>) {
    <span class="pl-s1">ctx</span><span class="pl-c1">-&gt;</span><span class="pl-c1">saved_data</span>[<span class="pl-s">"spatial_scale"</span>] <span class="pl-c1">=</span> <span class="pl-s1">spatial_scale</span>;
    <span class="pl-s1">ctx</span><span class="pl-c1">-&gt;</span><span class="pl-c1">saved_data</span>[<span class="pl-s">"pooled_height"</span>] <span class="pl-c1">=</span> <span class="pl-s1">pooled_height</span>;
    <span class="pl-s1">ctx</span><span class="pl-c1">-&gt;</span><span class="pl-c1">saved_data</span>[<span class="pl-s">"pooled_width"</span>] <span class="pl-c1">=</span> <span class="pl-s1">pooled_width</span>;
    <span class="pl-s1">ctx</span><span class="pl-c1">-&gt;</span><span class="pl-c1">saved_data</span>[<span class="pl-s">"sampling_ratio"</span>] <span class="pl-c1">=</span> <span class="pl-s1">sampling_ratio</span>;
    <span class="pl-s1">ctx</span><span class="pl-c1">-&gt;</span><span class="pl-c1">saved_data</span>[<span class="pl-s">"input_shape"</span>] <span class="pl-c1">=</span> <span class="pl-s1">input</span>.<span class="pl-en">sizes</span>();
    auto <span class="pl-smi">result</span> <span class="pl-c1">=</span> <span class="pl-en">PSROIAlign_forward</span>(
        <span class="pl-smi">input</span>,
        <span class="pl-smi">rois</span>,
        <span class="pl-smi">spatial_scale</span>,
        <span class="pl-smi">pooled_height</span>,
        <span class="pl-smi">pooled_width</span>,
        <span class="pl-smi">sampling_ratio</span>);
    auto <span class="pl-smi">output</span> <span class="pl-c1">=</span> <span class="pl-s1">std</span>::<span class="pl-s1">get</span><span class="pl-c1">&lt;</span><span class="pl-c1">0</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">result</span>);
    auto <span class="pl-smi">channel_mapping</span> <span class="pl-c1">=</span> <span class="pl-s1">std</span>::<span class="pl-s1">get</span><span class="pl-c1">&lt;</span><span class="pl-c1">1</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">result</span>);
    <span class="pl-s1">ctx</span><span class="pl-c1">-&gt;</span><span class="pl-en">save_for_backward</span>({<span class="pl-s1">rois</span>, <span class="pl-s1">channel_mapping</span>});
    <span class="pl-s1">ctx</span><span class="pl-c1">-&gt;</span><span class="pl-en">mark_non_differentiable</span>({<span class="pl-s1">channel_mapping</span>});
    <span class="pl-k">return</span> {<span class="pl-s1">output</span>, <span class="pl-s1">channel_mapping</span>};
  }

  <span class="pl-k">static</span> <span class="pl-smi">variable_list</span> <span class="pl-en">backward</span>(
      <span class="pl-smi">AutogradContext</span><span class="pl-c1">*</span> <span class="pl-s1">ctx</span>,
      <span class="pl-smi">variable_list</span> <span class="pl-s1">grad_output</span>) {
    <span class="pl-c">// Use data saved in forward</span>
    auto <span class="pl-smi">saved</span><span class="pl-s1"></span> <span class="pl-c1">=</span> <span class="pl-s1">ctx</span><span class="pl-c1">-&gt;</span><span class="pl-en">get_saved_variables</span>();
    auto <span class="pl-smi">rois</span> <span class="pl-c1">=</span> <span class="pl-s1">saved</span>[<span class="pl-c1">0</span>];
    auto <span class="pl-smi">channel_mapping</span> <span class="pl-c1">=</span> <span class="pl-s1">saved</span>[<span class="pl-c1">1</span>];
    auto <span class="pl-smi">input_shape</span><span class="pl-s1"></span> <span class="pl-c1">=</span> <span class="pl-s1">ctx</span><span class="pl-c1">-&gt;</span><span class="pl-c1">saved_data</span>[<span class="pl-s">"input_shape"</span>].<span class="pl-en">toIntList</span>();
    auto <span class="pl-smi">grad_in</span><span class="pl-s1"></span> <span class="pl-c1">=</span> <span class="pl-en">PSROIAlign_backward</span>(
        <span class="pl-s1">grad_output</span>[<span class="pl-c1">0</span>],
        <span class="pl-s1">rois</span>,
        <span class="pl-s1">channel_mapping</span>,
        <span class="pl-s1">ctx</span><span class="pl-c1">-&gt;</span><span class="pl-c1">saved_data</span>[<span class="pl-s">"spatial_scale"</span>].<span class="pl-en">toDouble</span>(),
        <span class="pl-s1">ctx</span><span class="pl-c1">-&gt;</span><span class="pl-c1">saved_data</span>[<span class="pl-s">"pooled_height"</span>].<span class="pl-en">toInt</span>(),
        <span class="pl-s1">ctx</span><span class="pl-c1">-&gt;</span><span class="pl-c1">saved_data</span>[<span class="pl-s">"pooled_width"</span>].<span class="pl-en">toInt</span>(),
        <span class="pl-s1">ctx</span><span class="pl-c1">-&gt;</span><span class="pl-c1">saved_data</span>[<span class="pl-s">"sampling_ratio"</span>].<span class="pl-en">toInt</span>(),
        <span class="pl-s1">input_shape</span>[<span class="pl-c1">0</span>],
        <span class="pl-s1">input_shape</span>[<span class="pl-c1">1</span>],
        <span class="pl-s1">input_shape</span>[<span class="pl-c1">2</span>],
        <span class="pl-s1">input_shape</span>[<span class="pl-c1">3</span>]);
    <span class="pl-k">return</span> {
        <span class="pl-s1">grad_in</span>, <span class="pl-en">Variable</span>(), <span class="pl-en">Variable</span>(), <span class="pl-en">Variable</span>(), <span class="pl-en">Variable</span>(), <span class="pl-en">Variable</span>()};
  }
};

<span class="pl-smi">std</span>::<span class="pl-s1">tuple</span><span class="pl-c1">&lt;</span><span class="pl-s1">Tensor</span>, <span class="pl-s1">Tensor</span><span class="pl-c1">&gt;</span> <span class="pl-en">ps_roi_align</span>(
    <span class="pl-s1">const</span> <span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">input</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">torch</span>::<span class="pl-s1">Tensor</span><span class="pl-c1">&amp;</span> <span class="pl-s1">rois</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">double</span> <span class="pl-s1">spatial_scale</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int64_t</span> <span class="pl-s1">pooled_height</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int64_t</span> <span class="pl-s1">pooled_width</span>,
    <span class="pl-s1">const</span> <span class="pl-s1">int64_t</span> <span class="pl-s1">sampling_ratio</span>) {
  auto <span class="pl-s1">result</span> <span class="pl-c1">=</span> <span class="pl-s1">PSROIAlignFunction</span>::<span class="pl-en">apply</span>(
      <span class="pl-s1">input</span>, <span class="pl-s1">rois</span>, <span class="pl-s1">spatial_scale</span>, <span class="pl-s1">pooled_height</span>, <span class="pl-s1">pooled_width</span>, <span class="pl-s1">sampling_ratio</span>);
  <span class="pl-k">return</span> <span class="pl-s1">std</span>::<span class="pl-s1">tuple</span><span class="pl-c1">&lt;</span><span class="pl-s1">Tensor</span>, <span class="pl-s1">Tensor</span><span class="pl-c1">&gt;</span>(<span class="pl-s1">result</span>[<span class="pl-c1">0</span>], <span class="pl-s1">result</span>[<span class="pl-c1">1</span>]);
}</pre></div>
</details>
<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/suo/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/suo">@suo</a></p>