在转ONNX的过程中出现了算子不支持的问题。grid_sampler是比较经典的算子不支持问题,在pytorch1.12之后已经支持了,而我这里要求的版本是torch=1.8.0,该了torch版本后,之后的推理平台又会有版本限制所以不能改。主要解决方式是通过算子替换来实现。
具体做法就是定位到你出错的代码位置,一般就是grid_sample函数的定义位置
然后用mmcv库中的bilinear_grid_sample()来替换grid_sample函数。二者是等价的,可以去torch官网看具体的api。
bilinear_grid_sample函数倒入方式
from mmcv.ops.point_sample import bilinear_grid_sample
如果这个时候还是报错,具体分析。比如报mmcv._ext没有,则先卸载然后重新下载mmcv-full
如果还是报错则重新写bilinear_grid_sample函数,具体代码如下
def bilinear_grid_sample(im, grid, align_corners=False):
"""Given an input and a flow-field grid, computes the output using input
values and pixel locations from grid. Supported only bilinear interpolation
method to sample the input pixels.
Args:
im (torch.Tensor): Input feature map, shape (N, C, H, W)
grid (torch.Tensor): Point coordinates, shape (N, Hg, Wg, 2)
align_corners (bool): If set to True, the extrema (-1 and 1) are
considered as referring to the center points of the input’s
corner pixels. If set to False, they are instead considered as
referring to the corner points of the input’s corner pixels,
making the sampling more resolution agnostic.
Returns:
torch.Tensor: A tensor with sampled points, shape (N, C, Hg, Wg)
"""
n, c, h, w = im.shape
gn, gh, gw, _ = grid.shape
assert n == gn
x = grid[:, :, :, 0]
y = grid[:, :, :, 1]
if align_corners:
x = ((x + 1) / 2) * (w - 1)
y = ((y + 1) / 2) * (h - 1)
else:
x = ((x + 1) * w - 1) / 2
y = ((y + 1) * h - 1) / 2
x = x.contiguous().view(n, -1)
y = y.contiguous().view(n, -1)
x0 = torch.floor(x).long()
y0 = torch.floor(y).long()
x1 = x0 + 1
y1 = y0 + 1
wa = ((x1 - x) * (y1 - y)).unsqueeze(1)
wb = ((x1 - x) * (y - y0)).unsqueeze(1)
wc = ((x - x0) * (y1 - y)).unsqueeze(1)
wd = ((x - x0) * (y - y0)).unsqueeze(1)
# Apply default for grid_sample function zero padding
im_padded = F.pad(im, pad=[1, 1, 1, 1], mode='constant', value=0)
padded_h = h + 2
padded_w = w + 2
# save points positions after padding
x0, x1, y0, y1 = x0 + 1, x1 + 1, y0 + 1, y1 + 1
# Clip coordinates to padded image size
x0 = torch.where(x0 < 0, torch.tensor(0), x0)
x0 = torch.where(x0 > padded_w - 1, torch.tensor(padded_w - 1), x0)
x1 = torch.where(x1 < 0, torch.tensor(0), x1)
x1 = torch.where(x1 > padded_w - 1, torch.tensor(padded_w - 1), x1)
y0 = torch.where(y0 < 0, torch.tensor(0), y0)
y0 = torch.where(y0 > padded_h - 1, torch.tensor(padded_h - 1), y0)
y1 = torch.where(y1 < 0, torch.tensor(0), y1)
y1 = torch.where(y1 > padded_h - 1, torch.tensor(padded_h - 1), y1)
#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#x0 = torch.where(x0 < 0, torch.tensor(0).to(device), x0)
#x0 = torch.where(x0 > padded_w - 1, torch.tensor(padded_w - 1).to(device), x0)
#x1 = torch.where(x1 < 0, torch.tensor(0).to(device), x1)
#x1 = torch.where(x1 > padded_w - 1, torch.tensor(padded_w - 1).to(device), x1)
#y0 = torch.where(y0 < 0, torch.tensor(0).to(device), y0)
#y0 = torch.where(y0 > padded_h - 1, torch.tensor(padded_h - 1).to(device), y0)
#y1 = torch.where(y1 < 0, torch.tensor(0).to(device), y1)
#y1 = torch.where(y1 > padded_h - 1, torch.tensor(padded_h - 1).to(device), y1)
im_padded = im_padded.view(n, c, -1)
x0_y0 = (x0 + y0 * padded_w).unsqueeze(1).expand(-1, c, -1)
x0_y1 = (x0 + y1 * padded_w).unsqueeze(1).expand(-1, c, -1)
x1_y0 = (x1 + y0 * padded_w).unsqueeze(1).expand(-1, c, -1)
x1_y1 = (x1 + y1 * padded_w).unsqueeze(1).expand(-1, c, -1)
Ia = torch.gather(im_padded, 2, x0_y0)
Ib = torch.gather(im_padded, 2, x0_y1)
Ic = torch.gather(im_padded, 2, x1_y0)
Id = torch.gather(im_padded, 2, x1_y1)
return (Ia * wa + Ib * wb + Ic * wc + Id * wd).reshape(n, c, gh, gw)
可能会出现下面情况,把 # Clip coordinates to padded image size部分注释的取消,上面的八行代码删掉
如果出现下面问题,在刚才grid_sample函数那里把reshape改改
基本上就差不多了。
参考
https://blog.youkuaiyun.com/Lizongming_/article/details/130065380?utm_medium=distribute.pc_relevant.none-task-blog-2~default~baidujs_baidulandingword~default-0-130065380-blog-125559944.235^v38^pc_relevant_anti_vip_base&spm=1001.2101.3001.4242.1&utm_relevant_index=3#commentBox
https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html