torch.rfft相关的一个小问题

[chenli199]

因为新版torch的原因，在网络上查了一些新旧更替的资料，发现对于图像数据，torch.rfft(img, 3, onesided=False)这里的dim给的都是2，不是3，请问是有什么影响吗？

[chenli199]

[Algolzw]

@chenli199

• 在我们代码里因为是对每个通道单独做的，所以2或3维应该都可以。

• 我也没有转到过新版本的rfft实验过代码，不过ifft2出来的应该是复数，不知道你取实部没有(deblur_f.real)。具体的转化我可能要花点时间测一下。

[RC-Qiao]

请问这个问题解决了么， 我想用3090的跑但pytorch新版本的rfft和旧版有问题

[LeoChen12]

您好 请问这个rfft的版本问题解决了吗 新旧版本无法匹配

[Algolzw]

@LeoChen12 @hdcxqrc 可以尝试以下代码

def normkernel_to_downkernel(rescaled_blur_hr, rescaled_hr, ksize, eps=1e-10):
    blur_img = torch.fft.fftn(rescaled_blur_hr, dim=(-3, -2, -1))
    blur_img = torch.stack((blur_img.real, blur_img.imag), -1)
    img = torch.fft.fftn(rescaled_hr, dim=(-3, -2, -1))
    img = torch.stack((img.real, img.imag), -1)

    denominator = img[:, :, :, :, 0] * img[:, :, :, :, 0] \
                  + img[:, :, :, :, 1] * img[:, :, :, :, 1] + eps

    inv_denominator = torch.zeros_like(img)
    inv_denominator[:, :, :, :, 0] = img[:, :, :, :, 0] / denominator
    inv_denominator[:, :, :, :, 1] = -img[:, :, :, :, 1] / denominator

    kernel = torch.zeros_like(blur_img).cuda()
    kernel[:, :, :, :, 0] = inv_denominator[:, :, :, :, 0] * blur_img[:, :, :, :, 0] \
                            - inv_denominator[:, :, :, :, 1] * blur_img[:, :, :, :, 1]
    kernel[:, :, :, :, 1] = inv_denominator[:, :, :, :, 0] * blur_img[:, :, :, :, 1] \
                            + inv_denominator[:, :, :, :, 1] * blur_img[:, :, :, :, 0]
    ker = convert_otf2psf(kernel, ksize)
    return ker

def convert_otf2psf(otf, size):
    ker = torch.zeros(size)
    psf = torch.fft.ifft2(torch.complex(otf[..., 0], otf[..., 1]), dim=(-3, -2, -1))
    # circularly shift
    ksize = size[-1]
    centre = ksize//2 + 1

    ker[:, :, (centre-1):, (centre-1):] = psf[:, :, :centre, :centre]#.mean(dim=1, keepdim=True)
    ker[:, :, (centre-1):, :(centre-1)] = psf[:, :, :centre, -(centre-1):]#.mean(dim=1, keepdim=True)
    ker[:, :, :(centre-1), (centre-1):] = psf[:, :, -(centre-1):, :centre]#.mean(dim=1, keepdim=True)
    ker[:, :, :(centre-1), :(centre-1)] = psf[:, :, -(centre-1):, -(centre-1):]#.mean(dim=1, keepdim=True)

    return ker

[LeoChen12]

感谢您的帮助！