Changing the shap of the output of Segformer

Models
1

I am using the Segformer model for image segmentation on the Kvasir-SEG dataset and my input has shape (32, 3, 225, 225). I want the output of the model to have the same shape, but when I run the following code I get an output of shape (32, 256, 7, 7). How can I modify my code to get the output to have the same shape as the input?

from transformers import SegformerModel, SegformerConfig

import torch

# Initializing a SegFormer nvidia/segformer-b0-finetuned-ade-512-512 style configuration

configuration = SegformerConfig()

# Initializing a model from the nvidia/segformer-b0-finetuned-ade-512-512 style configuration

model = SegformerModel(configuration)

# Generate some random input data

value = torch.rand(32, 3, 224, 224)

# Forward pass through the model

output = model(value, return_dict=False)[0]

# Print the output shape

print(output.shape) # should be (32, 3, 224, 224)
1 Like
2

I am also interested in this. Segmentation models are kinda useless if I am not able to recreate the ground truth with the exact same shape.