pytorch jit 内存泄漏（Pytorch中的grid_sample算子功能解析）

         pytorch中的grid_sample是一种特殊的采样算法    
    
    
。

调用接口为：

torch.nn.functional.grid_sample(input,grid,mode=bilinear,padding_mode=zeros,align_corners=None) 


     


     


   。

         input参数是输入特征图tensor
    
    
   ，也就是特征图     


     


     

，可以是四维或者五维张量 




 




 




，以四维形式为例(N,C,Hin,Win)

   

   

  ，N可以理解为Batch_size,C可以理解为通道数     

     

     
，Hin和Win也就是特征图高和宽

 




 




 

。

         grid包含输出特征图特征图的格网大小以及每个格网对应到输入特征图的采样点位  




  




  




，对应四维input


  


  


 ，其张量形式为(N,Hout,Wout,2)     
     
     ，其中最后一维大小必须为2   




   




   




，如果输入为五维张量



 



 



，那么最后一维大小必须为3   

   

   

。为什么最后一维必须为2或者3？因为grid的最后一个维度实际上代表一个坐标(x,y)或者(xy,z)               ，对应到输入特征图的二维或三维特征图的坐标维度    




    




    



，xy取值范围一般为[-1,1]














，该范围映射到输入特征图的全图  

     

     

  。

         mode为选择采样方法              ，有三种内插算法可选     



     



     


，分别是bilinear双线性差值    
    
    
、nearest最邻近插值 


     


     


   、bicubic 双三次插值


  




  




  
。

         padding_mode为填充模式














，即当(x,y)取值超过输入特征图采样范围
    
    
   ，返回一个特定值     


     


     

，有zeros 

 




 




 

、 border    

   

   

、 reflection三种可选 




 




 




，一般用zero  


  


  


。

         align_corners为bool类型

   

   

  ，指设定特征图坐标与特征值对应方式     

     

     
，设定为TRUE时  




  




  




，特征值位于像素中心   
     
     
 。

         要理解grid_sample是如何工作的


  


  


 ，最好就是进行简单的复现



   




   




   。假设输入shape为(N,C,H,W)     
     
     ，grid的shape设定为(N,H,W,2)   




   




   




，以双线性差值为例进行处理 



 



 



。首先根据input和grid设定



 



 



，输出特征图tensor的shape为(N,C,H,W)               ，输出特征图上每一个cell上的值由grid最后一维(x,y)确定              。那么如何计算输出tensor上每一个点的值？首先    




    




    



，通过(x,y)找到输入特征图上的采样位置














，由于xy取值范围为[-1,1]              ，为了便于计算     



     



     


，先将xy取值范围调整为[0,1]




    




    




    。通过(w-1)*(x+1)/2  

     

     

  、(wh-1)*(y+1)/2将xy映射为输入特征图的具体坐标位置













。将xy映射到特征图实际坐标后














，取该坐标附近四个角点特征值
    
    
   ，通过四个特征值坐标与采样点坐标相对关系进行双线性插值     


     


     

，得到采样点的值               。

注意：xy映射后的坐标可能是输入特征图上任意位置



     



     



    。假设输出特征图上(2,2)坐标位置上的值采样位置可能为输入特征图上(3,4)位置 




 




 




，xy越小越靠近输入特征图左上角

   

   

  ，越大则越靠近右下角













。

         基于上面的思路     

     

     
，可以进行一个简单的自定义实现    
    
    
。根据指定shape生成input和grid  




  




  




，使用pytorch中的grid_sample算子生成output 


     


     


   。之后取grid中的第一个位置中的xy


  


  


 ，根据xy从input中通过双线性插值计算出output第一个位置的值

 




 




 

。

import torch import numpy as np def grid_sample(input, grid): N, C, H_in, W_in = input.shape N, H_out, W_out, _ = grid.shape output = np.random.random((N,C,H,W)) for i in range(N): for j in range(C): for k in range(H_out): for l in range(W_out): param = [0.0, 0.0] param[0] = (W_in - 1) * (grid[i][k][l][0] + 1) / 2 param[1] = (H_in - 1) * (grid[i][k][l][1] + 1) / 2 x0 = int(param[0]) x1 = x0 + 1 y0 = int(param[1]) y1 = y0 + 1 param[0] -= x0 param[1] -= y0 left_top = input[i][j][y0][x0] * (1 - param[0]) * (1 - param[1]) left_bottom = input[i][j][y1][x0] * (1 - param[0]) * param[1] right_top = input[i][j][y0][x1] * param[0] * (1 - param[1]) right_bottom = input[i][j][y1][x1] * param[0] * param[1] result = left_bottom + left_top + right_bottom + right_top output[i][j][k][l] = result return output N, C, H, W = 1, 1, 4, 4 input = np.random.random((N,C,H,W)) grid = np.random.random((N,H,W,2)) out = grid_sample(input, grid) print(f自定义实现输出结果：\n{out}) input = torch.from_numpy(input) grid = torch.from_numpy(grid) output = torch.nn.functional.grid_sample(input,grid,mode=bilinear, padding_mode=zeros,align_corners=True) print(fgrid_sample输出结果：\n{output})

运行结果：

         从输出结果上看     
     
     ，与pytorch基本一致   




   




   




，由于仅仅做简单验证



 



 



，这里没有对超出[-1,1]范围的xy值做处理               ，只能处理四维input    




    




    



，五维input的实现思路与这里基本一致   

   

   

。

        考虑到(x,y)取值范围可能越界














，pytorch中的padding_mode设置就是对(x,y)落在输入特征图外边缘情况进行处理              ，一般设置zero     



     



     


，也就是对靠近输入特征图范围以外的采样点进行0填充














，如果不进行处理显然会造成索引越界  

     

     

  。要解决(x,y)越界问题
    
    
   ，可以进行如下修改：

import torch import numpy as np def grid_sample(input, grid): N, C, H_in, W_in = input.shape N, H_out, W_out, _ = grid.shape output = np.random.random((N, C, H_out, W_out)) for i in range(N): for j in range(C): for k in range(H_out): for l in range(W_out): x, y = grid[i][k][l][0], grid[i][k][l][1] param = [0.0, 0.0] param[0] = (W_in - 1) * (x + 1) / 2 param[1] = (H_in - 1) * (y + 1) / 2 x1 = int(param[0] + 1) x0 = x1 - 1 y1 = int(param[1] + 1) y0 = y1 - 1 param[0] = abs(param[0] - x0) param[1] = abs(param[1] - y0) left_top_value, left_bottom_value, right_top_value, right_bottom_value = 0, 0, 0, 0 if 0 <= x0 < W_in and 0 <= y0 < H_in: left_top_value = input[i][j][y0][x0] if 0 <= x1 < W_in and 0 <= y0 < H_in: right_top_value = input[i][j][y0][x1] if 0 <= x0 < W_in and 0 <= y1 < H_in: left_bottom_value = input[i][j][y1][x0] if 0 <= x1 < W_in and 0 <= y1 < H_in: right_bottom_value = input[i][j][y1][x1] left_top = left_top_value * (1 - param[0]) * (1 - param[1]) left_bottom = left_bottom_value * (1 - param[0]) * param[1] right_top = right_top_value * param[0] * (1 - param[1]) right_bottom = right_bottom_value * param[0] * param[1] result = left_bottom + left_top + right_bottom + right_top output[i][j][k][l] = result return output N, C, H_in, W_in = 1, 1, 4, 4 H_out, W_out = 4, 4 input = np.random.random((N, C, H_in, W_in)) grid = np.random.random((N, H_out, W_out, 2)) grid[0][0][0] = [-1.2, 1.3] out = grid_sample(input, grid) print(f自定义实现输出结果：\n{out}) input = torch.from_numpy(input) grid = torch.from_numpy(grid) output = torch.nn.functional.grid_sample(input, grid, mode=bilinear, padding_mode=zeros, align_corners=True) print(fgrid_sample输出结果：\n{output})

     测试结果：