diff --git a/segmentation_models_pytorch/decoders/segformer/model.py b/segmentation_models_pytorch/decoders/segformer/model.py index 65d7e8fa..03deeeef 100644 --- a/segmentation_models_pytorch/decoders/segformer/model.py +++ b/segmentation_models_pytorch/decoders/segformer/model.py @@ -29,6 +29,7 @@ class Segformer(SegmentationModel): activation: An activation function to apply after the final convolution layer. Available options are **"sigmoid"**, **"softmax"**, **"logsoftmax"**, **"tanh"**, **"identity"**, **callable** and **None**. Default is **None**. + upsampling: A number to upsample the output of the model, default is 4 (same size as input) aux_params: Dictionary with parameters of the auxiliary output (classification head). Auxiliary output is build on top of encoder if **aux_params** is not **None** (default). Supported params: - classes (int): A number of classes @@ -56,6 +57,7 @@ def __init__( in_channels: int = 3, classes: int = 1, activation: Optional[Union[str, Callable]] = None, + upsampling: int = 4, aux_params: Optional[dict] = None, **kwargs: dict[str, Any], ): @@ -80,7 +82,7 @@ def __init__( out_channels=classes, activation=activation, kernel_size=1, - upsampling=4, + upsampling=upsampling, ) if aux_params is not None: