CoTNet
CoTNet copied to clipboard
Published
20 hours ago •
JDAI-CV
JDAI-CV
export onnx error
cuda11.1 + pytorch1.8 + cupy-cuda111, how to fix it, thanks a lot.
根据test_aggregation_zeropad函数将aggregation_zeropad用torch api 实现，以下仅供参考：
def custom_aggregation_zeropad(input, weight, kernel_size=3, stride=1, padding=0, dilation=1):
"""
自定义，用torch api实现
"""
head_num = 2
n, c_x, c_w, in_height, in_width = input.shape[0], input.shape[1], weight.shape[2], input.shape[-2], input.shape[-1]
out_height = int((in_height + 2 * padding - (dilation * (kernel_size - 1) + 1)) / stride + 1)
out_width = int((in_width + 2 * padding - (dilation * (kernel_size - 1) + 1)) / stride + 1)
unfold_j = torch.nn.Unfold(kernel_size=kernel_size, dilation=dilation, padding=padding, stride=stride)
x2 = unfold_j(input).view(n, c_x // c_w, c_w, pow(kernel_size, 2), out_height, out_width)
y2 = (weight.unsqueeze(2) * x2.unsqueeze(1)).sum(-3).view(n, head_num * c_x, out_height, out_width)
return y2
