Sensitivity of SS2D Block to Input Image Feature Size

[Guaishou74851]

Description:

I have trained an EDSR-style network with 20 residual blocks (RB). The baseline network has 20 RBs, while the experimental variant inserts one MySS2D block in the middle, resulting in the structure: 10RBs + 1MySS2D + 10RBs.

Here are the details of my observations:

1. Training Details:

   • Patch Size: 128x128

2. Testing Observations:

   • Using 128x128 Blocks: When I divide the input image feature into 128x128 blocks before passing through the SS2D block, the PSNR is normal.
   • Using Original 512x512 Image: When I directly input the original image of 512x512 size, the PSNR is significantly lower, with around a 5dB drop compared to the baseline network.

This leads me to suspect that the SS2D block might be sensitive to the input image block size and possibly overfitted to the patch size used during training.

Code Implementation:

from mambair_arch import SS2D

class MySS2D(nn.Module):
    def __init__(self, C):
        super().__init__()
        self.body = SS2D(C)
        self.s = nn.Parameter(torch.tensor([0.0]))
    
    def forward(self, x):
        B, C, H, W = x.shape
        return x + self.s * self.body(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)

Questions:

1. Is the SS2D block sensitive to the input image feature size?
2. Is there a potential mistake in my implementation that could be causing this issue?
3. Any insights or instructions on how to address this issue would be greatly appreciated.

Thank you very much for your assistance!

[csguoh]

I have replied through e-mail.