pytorch实现图像投影变换

import cv2
import torch


def cpu_remap(numpy_img,mapx,mapy):
    return cv2.remap(numpy_img,mapx,mapy,cv2.INTER_LINEAR)

def gpu_remap(numpy_img,map_tensor):
    '''
    numpy_img:原始图像格式为ndarray
    map_tensor:[N,H,W,C]用于grid_sample的map参数，需要规制到-1到1
    '''
    # 准备图像数据
    img_tensor = torch.from_numpy(numpy_img).contiguous().cuda(non_blocking=True)
    img_tensor = img_tensor.permute(2,0,1).unsqueeze(0).cuda().float()
    res = torch.nn.functional.grid_sample(img_tensor,map_tensor,
                                          mode='bilinear',
                                          padding_mode='zeros',
                                          align_corners=None)
    res = res.char()
    res = res[0].permute(1,2,0)
    res = res.cpu()
    res = res.numpy()
    res = np.uint8(res)
    return res
@profile
def main():
   #已知单应矩阵（投影矩阵M）
   M=np.asarray([[-5.23249213e+00, -2.83428439e+01,  2.74163372e+03],
       [-2.71994329e+00, -2.52929752e+01,  1.13544900e+04],
       [-1.38531350e-04, -1.40379841e-02,  1.00000000e+00]])
   
   
   img_revise = cv2.imread('206_revise.jpg') # 无畸变的图像
   img_perspect = cv2.warpPerspective(img_revise, M, (2800, 1500), borderValue=0)
   cv2.imwrite('warpPerspective.jpg',img_perspect)

   x, y = np.meshgrid(np.arange(2800), np.arange(1500))
   grid = np.vstack([x.flatten(),y.flatten(),np.ones(y.flatten().shape)])
   grid_trans = np.linalg.inv(M)@grid
   grid_trans = grid_trans/grid_trans[2]
   grid_trans = grid_trans[:2]
   mapx = np.float32(np.reshape(grid_trans[0],[1500,2800]))#/1920*2-1
   mapy = np.float32(np.reshape(grid_trans[1],[1500,2800]))#/1080*2-1
   print(mapx.shape)
   img_perspcet = cpu_remap(img_revise,mapx,mapy)
   cv2.imwrite('cpu_remap.jpg',img_perspect)

   grid_trans_x = torch.from_numpy(mapx).unsqueeze(2)/1920*2-1
   grid_trans_y = torch.from_numpy(mapy).unsqueeze(2)/1080*2-1
   map_tensor = torch.cat([grid_trans_x,grid_trans_y],2).unsqueeze(0).float().cuda()
   img_perspect = gpu_remap(img_revise,map_tensor)
   cv2.imwrite('gpu_remap.jpg',img_perspect)
if __name__=='__main__':
    main()

还是那个问题，数据在cpu->GPU拷贝时会慢，单GPU的计算耗时很短