Models ====== DeepTrees' crown-delineation model is made of two sub-networks -- a segmentation model (predicts the tree-cover mask and crown outlines) and a distance model (predicts the distance transform used to separate touching crowns). Both are built with `segmentation_models_pytorch `_ (SMP) and are fully configurable: you can switch the underlying architecture and encoder backbone through the ``model.architecture`` and ``model.backbone`` keys in your config YAML, without touching any code. Supported architectures ------------------------ Set via ``model.architecture``. Every architecture below is available to both the segmentation and distance sub-networks: .. list-table:: :header-rows: 1 * - ``architecture`` - Notes * - ``Unet`` - Default. Classic encoder-decoder with skip connections. Good general-purpose choice. * - ``Unet++`` - Unet with nested, densely-connected skip pathways. * - ``Linknet`` - Lightweight encoder-decoder, cheaper than Unet. * - ``FPN`` - Feature Pyramid Network decoder. * - ``PSPNet`` - Pyramid Scene Parsing decoder. * - ``PAN`` - Pyramid Attention Network decoder. * - ``DeepLabV3`` - Atrous/dilated convolution decoder. * - ``DeepLabV3+`` - DeepLabV3 with an added decoder for sharper boundaries. * - ``Segformer`` - Transformer-based (SegFormer). Requires a Mix Vision Transformer backbone -- see below. Backbones --------- Set via ``model.backbone``. For every architecture except ``Segformer``, this can be any encoder supported by SMP -- e.g. ``resnet18``, ``resnet34``, ``resnet50``, ``efficientnet-b0``, and many more. See SMP's own `encoder list `_ for the full set. ``Segformer`` is the one exception: it requires a **Mix Vision Transformer** backbone -- ``mit_b0`` through ``mit_b5`` (``mit_b0`` is the smallest/fastest, ``mit_b5`` the largest). Passing a CNN backbone (e.g. ``resnet18``) with ``architecture: Segformer`` will fail, since SMP's Segformer decoder is only wired up for MiT-style encoder outputs. Example: switching architectures --------------------------------- The default (``Unet`` + ``resnet18``), as used in the shipped predict/train configs: .. code-block:: yaml model: in_channels: 5 architecture: Unet backbone: resnet18 Switching to a ``Segformer`` model instead only requires changing these two lines: .. code-block:: yaml model: in_channels: 5 architecture: Segformer backbone: mit_b0 Everything else in the config (data, augmentation, postprocessing, trainer settings) stays the same -- ``architecture``/``backbone`` are the only knobs that change which model gets built. Encoder pretrained weights --------------------------- Independently of which architecture/backbone you pick, ``model.encoder_weights`` (train configs) controls whether the encoder starts from SMP's ImageNet-pretrained weights (``imagenet``) or random initialization (``null``). This is unrelated to loading a full DeepTrees checkpoint (below) -- it only affects the *encoder*, and only matters when training a new model, not when loading an already-trained one. Loading a pretrained DeepTrees checkpoint ------------------------------------------- This is the part that catches people out: **the config's architecture/backbone is not read from the checkpoint file.** It has to be set correctly *before* the checkpoint is loaded, because the predict/inference code builds the model from ``model.architecture``/``model.backbone`` first, and only then loads weights into that already-built structure. Nothing inspects the checkpoint to figure out what produced it. For prediction (``config/predict/*.yaml``), the checkpoint is set via: .. code-block:: yaml model: in_channels: 5 architecture: Segformer # must match how the checkpoint was trained backbone: mit_b0 # must match how the checkpoint was trained download_pretrained_model: False pretrained_model_path: ./pretrained_models pretrained_model_name: my_segformer_checkpoint_jitted.pt For training/fine-tuning (``config/train/*.yaml``), the equivalent is the top-level ``pretrained`` block (``pretrained.path`` / ``pretrained.model``). If ``architecture``/``backbone`` don't match what the checkpoint actually is, loading fails loudly and immediately with a ``RuntimeError`` listing missing/mismatched state dict keys (e.g. a ResNet checkpoint loaded against a ``Segformer``-configured model will complain about missing keys like ``encoder.conv1.weight``, which only exist in a ResNet encoder). This is a safe failure mode -- it will not silently produce a broken or nonsensical model -- but there's no upfront check; you only find out at load time. ``in_channels`` must match the checkpoint too, for the same reason. .. note:: Different pretrained checkpoints are not interchangeable just because they're both ``.pt`` files. A ``Unet``/``resnet18`` checkpoint only loads into a ``Unet``/``resnet18``-configured model, and likewise for ``Segformer``/``mit_b0``. Keep track of which architecture/backbone each of your checkpoints was trained with. Known limitations ------------------ - **Segformer training does not currently work on Apple's MPS backend.** It hits a PyTorch/MPS operator gap in a transformer reshape operation (``view size is not compatible with input tensor's size and stride``). Training falls back to CPU, which works correctly but without GPU acceleration on Mac. CUDA GPUs are unaffected.