Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

SegFormer Training extremely slow #996

Open
omarequalmars opened this issue Dec 4, 2024 · 2 comments · May be fixed by #998
Open

SegFormer Training extremely slow #996

omarequalmars opened this issue Dec 4, 2024 · 2 comments · May be fixed by #998

Comments

@omarequalmars
Copy link

I've been training a series of models implemented by this package for a while, all using 'tu-mobilevit_xxs' as an encoder. However I noticed that the latest addition, segformer, is extremely slow in training compared to the others despite being reported to have the same number of parameters by Pytorch lightning. Here is a visualization from Tensorboard:

image

And here is the parameter count reported by Lightning:

image

It's much slower than the rest in training, and considerably slows down my laptop. Why does it take so much resources despite being approximately the same size as the other models? Here are my architecture hyperparameters for reference, for every model I trained:

    h_params_Unet = {
        'encoder': encoder,
        'depth': 5,
        'Bnorm': True,
        'Attn': 'scse',
        'channels': (256, 128, 64, 32, 16),
        'lr_init': 1e-3,
        'weight_decay': 0,
        'temperature': 2,
        'Arch': "Unet"
    }

    h_params_DLV3plus = {
        'encoder': encoder,
        'depth': 5,
        'encoder_output_stride': 16,
        'channels': 512,
        'decoder_atrous_rates': (16, 32, 128),
        'lr_init': 1e-3,
        'weight_decay': 0,
        'temperature': 2,
        'Arch': "DLV3+"
    }

    h_params_MANet = {
        'encoder': encoder,               # Backbone encoder
        '
```depth': 5,                       # Number of encoder stages
        'Bnorm': True,                    # Use batch normalization in the decoder
        'channels': (256, 128, 64, 32, 16),  # Number of channels in each decoder layer
        'decoder_pab_channels': 64,
        'lr_init': 1e-3,                  # Initial learning rate
        'weight_decay': 0,                # Weight decay for regularization
        'temperature': 2,                 # Temperature for distillation
        'Arch': "MANet"
    }

    h_params_PAN = {
        'encoder': encoder,               # Backbone encoder
        'encoder_output_stride': 16,
        'channels': 512,                  # Number of channels in each decoder layer
        'lr_init': 1e-3,                  # Initial learning rate
        'weight_decay': 0,                # Weight decay for regularization
        'temperature': 2,                 # Temperature for distillation
        'Arch': "PAN"
    }

    h_params_Segformer = {
    'encoder': encoder,               # Backbone encoder
    'depth': 5,                       # Number of encoder stages
    'channels': 512,  # Number of channels in each decoder layer
    'lr_init': 1e-3,                  # Initial learning rate
    'weight_decay': 0,                # Weight decay for regularization
    'temperature': 2,                 # Temperature for distillation
    'Arch': "Segformer"
}


All of them have the same number of channels and depth where applicable. What am I doing wrong?
@qubvel
Copy link
Collaborator

qubvel commented Dec 4, 2024

Hey @omarequalmars, thanks for reporting the issue! It might be due to it's high-resolution features in the head.. but there might be some profiling needed to find the reason.

cc @brianhou0208 maybe you have some insights

@brianhou0208 brianhou0208 linked a pull request Dec 5, 2024 that will close this issue
@brianhou0208
Copy link
Contributor

Hi @omarequalmars ,

I tested your hyperparameter settings and resolved performance issue. Please check #998

Given that you're using a lightweight backbone encoder, I recommend reducing the decoder_channels to further optimize inference speed, as highlighted in the original paper.

Regarding input resolution, since it wasn't specified, smaller resolutions are unlikely to cause significant speed differences. However, larger resolutions could have a notable impact on performance.

If you're exploring transformer-style semantic segmentation models, aside from SegFormer, you may also consider TopFormer and SeaFormer, which are designed for efficiency as well.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging a pull request may close this issue.

3 participants