segmentation_models_pytorch_example/base/__init__.py

from .model import SegmentationModel

from .modules import (
    Conv2dReLU,
    Attention,
)

from .heads import (
    SegmentationHead,
    ClassificationHead,
)