fft cannot export to onnx, can we use Gaussian convolution to filter background?

[PonyPinkPie]

Hit, @ChengShiest

Thank you very much for open-sourcing such an outstanding work, and the improvement in performance metrics is tremendous.

However, the problem still persists. The operators in the FFT series cannot be exported to ONNX and TensorRT. I have considered an alternative solution, which is to use Gaussian convolution as a substitute.

I believe you have also considered this solution before. Is this feasible?

Before attempting to perform ImageNet pre-training, I would like to inquire, perhaps you have already conducted related experiments.

Thank you for any reply :)

class dense_vit(VisionTransformer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.gaussian_conv = nn.Conv1d(
            in_channels=768,
            out_channels=768,
            kernel_size=7,
            padding=3,
            groups=768, 
            bias=False
        )
        
        self._init_gaussian_weight()

    def gaussian_kernel_1d(self, kernel_size, sigma):

        x = torch.arange(-kernel_size // 2 + 1, kernel_size // 2 + 1, dtype=torch.float32)
        gauss = torch.exp(-(x ** 2) / (2 * sigma ** 2))
        gauss = gauss / gauss.sum()
        return gauss

    def _init_gaussian_weight(self):

        kernel = self.gaussian_kernel_1d(kernel_size=7, sigma=1.0)
        kernel = kernel.view(1, 1, 7).repeat(768, 1, 1)  # [groups, 1, k]
        self.gaussian_conv.weight.data = kernel
        self.gaussian_conv.weight.requires_grad = False  

    def forward(self, x: torch.Tensor):
        x = self._process_input(x)
        n = x.shape[0]

        batch_class_token = self.class_token.expand(n, -1, -1)
        x = torch.cat([batch_class_token, x], dim=1)
        x = self.encoder(x)

        x_detach = x[:, 1:]

        x_in = x_detach.transpose(1, 2)  # [B, 768, N]

        x_smooth = self.gaussian_conv(x_in)  # [B,768,N]

        x_smooth = x_smooth.transpose(1, 2)

        diff = x_detach / (torch.abs(x_detach - x_smooth) + 1e-6)

        _, indices = torch.topk(diff, k=1, dim=1, largest=True)

        sel_p = torch.gather(x_detach, 1, indices)

        cls_token = torch.mean(sel_p, dim=1)

        return cls_token, x_detach

[ChengShiest]

Thank you very much for your interest in our work and for the detailed discussion.

I am not very familiar with ONNX and TensorRT deployment, so I may not be able to give a definitive answer from the engineering/export side. But regarding the method you mentioned, my intuition is that replacing the FFT-based smoothing/filtering operation with a Gaussian convolution should be feasible.

Thanks again for your thoughtful suggestion :)