Skip to main content
SegmentationTrainConfig extends TrainConfig with additional parameters that control mask prediction losses. It is used automatically when you call model.train() on any segmentation variant (RFDETRSegNano, RFDETRSegSmall, RFDETRSegMedium, RFDETRSegLarge, RFDETRSegXLarge, RFDETRSeg2XLarge). You do not need to instantiate this class directly — pass parameters as keyword arguments to model.train(): 
from rfdetr import RFDETRSegMedium

model = RFDETRSegMedium()
model.train(
    dataset_dir="dataset/",
    epochs=50,
    batch_size=4,
    mask_ce_loss_coef=5.0,
    mask_dice_loss_coef=5.0,
)
SegmentationTrainConfig inherits all fields from TrainConfig. See TrainConfig for the complete list of inherited parameters (learning rate, batch size, logging, multi-GPU, etc.). Only the additional and overridden fields are documented on this page.

​
Segmentation-specific fields

​
mask_point_sample_ratio
int
default:"16"
Number of points sampled per mask for the point-based cross-entropy loss. A higher value increases loss quality at the cost of more memory. Corresponds to the num_points in the PointRend-style mask loss.
​
mask_ce_loss_coef
float
default:"5.0"
Weight applied to the binary cross-entropy component of the mask loss.
​
mask_dice_loss_coef
float
default:"5.0"
Weight applied to the Dice loss component of the mask loss.

​
Overridden defaults

The following TrainConfig fields have different defaults in SegmentationTrainConfig:
​
cls_loss_coef
float
default:"5.0"
Weight applied to the classification loss. Overrides the TrainConfig default of 1.0 to match the higher mask loss coefficients.
​
segmentation_head
boolean
default:"true"
Always True for segmentation models. This field is fixed and cannot be overridden.

​
Full usage example

from rfdetr import RFDETRSegMedium

model = RFDETRSegMedium()

model.train(
    # Required
    dataset_dir="dataset/",
    # Core
    epochs=75,
    output_dir="runs/seg-medium",
    # Batch & memory
    batch_size="auto",
    grad_accum_steps=4,
    # Learning rate
    lr=1e-4,
    lr_encoder=1.5e-4,
    lr_scheduler="cosine",
    # Segmentation loss weights
    mask_ce_loss_coef=5.0,
    mask_dice_loss_coef=5.0,
    cls_loss_coef=5.0,
    # Logging
    tensorboard=True,
    wandb=False,
)
Run inference after training to see the segmentation masks: 
detections = model.predict("image.jpg", threshold=0.5)

# detections.mask is an np.ndarray of shape (N, H, W)
for i, mask in enumerate(detections.mask):
    print(f"Object {i}: class={detections.class_id[i]}, confidence={detections.confidence[i]:.2f}")

TrainConfig

Complete reference for all inherited training parameters.

Model variants

All available segmentation model classes and their default configs.

Run segmentation

Guide to running inference with segmentation models.

Training overview

End-to-end training guide for custom datasets.

Build docs developers (and LLMs) love

Get started for freeTalk to us