[FEATURE] Support variable input size of maxvit and coatnet

[abebe9849]

Will it be possible in the future to support variable input sizes for maxvit and coatnet?

I am experimenting with adapting various models of timm to self-supervised learning such as DINO and iBOT. When learning these, it is necessary to accept inputs of (224,224,3) and (96,96,3) at the same time.

The model here seems to support variable sizes. If you're not going to break the implementation, I'd love to see you add support for variable sizes.

[tlpss]

@abebe9849 I was also looking into the flexible input sizes of MaxViT (although for finetuning on a different size)

AFAIK, the MaxVit implementation accepts variable input sizes but the size must be divisible by the number of partitions (P) and grids (G) for each (downsampled) image size. Since the configurations in timm and in the paper downsample up to 32 times, this implies that your input size needs to be a multiple of 32*P/G.

So i'm thinking that in your case you need to pad/crop/resize your images to 96 & 192 or 112 & 224 (but then you have no pretrained weights) or to multiples of 224 (and use pretrained weights). In the first cases you would need to specify the image size for the timm model (which will be used in timm to find the appropriate P/G ).

This can be done as follows:

maxvit = timm.create_model("maxvit_tiny_rw_224", pretrained=False,img_size = 96)
inputs= [torch.zeros((1,3,96,96)), torch.zeros((1,3,192,192))]
for x in inputs:
    y= maxvit(x)
    print(y.shape)

resulting in:

torch.Size([1, 1000])
torch.Size([1, 1000])

From section 3 of the accompanying notebook in the original repo, it seems that the features are scale-invariant enough to handle such large differences in input size so that should be fine.