Fine-tuning
Train on your own image/mask pairs with the HuggingFace Trainer.
nobg provides the model, the processor and the loss; because the models plug into the HuggingFace
Trainer, you get its training loop, checkpointing and evaluation for free.
The contract is small:
- the processor turns
segmentation_mapsinto binarizedlabelsof shape(B, 1, H, W) forward(pixel_values, labels=...)returnslossalongsidelogitsmodel.criterionis the function that computes it, and it is swappable
BiRefNet
from nobg import AutoModel, AutoProcessor
from transformers import Trainer, TrainingArguments
model = AutoModel.from_pretrained("feyninc/FeyNobg")
processor = AutoProcessor.from_pretrained("feyninc/FeyNobg")
def collate(examples):
batch = processor(
images=[example["image"] for example in examples],
segmentation_maps=[
example["mask"].convert("L") for example in examples
],
return_tensors="pt",
)
return {
"pixel_values": batch["pixel_values"],
"labels": batch["labels"],
}
trainer = Trainer(
model=model,
args=TrainingArguments(
output_dir="outputs",
learning_rate=2e-5,
),
train_dataset=dataset,
data_collator=collate,
)
trainer.train()dataset is any torch.utils.data.Dataset or datasets.Dataset whose items carry an image and a
mask — the collator is what adapts your field names, so nothing else has to change.
segmentation_maps are resized and rescaled like the images but never normalized, then binarized at
0.5 — BiRefNet trains against a hard mask even though it predicts a soft one. The channel dimension is
kept, so labels arrive as (B, 1, H, W).
The loss is birefnet_loss: a multi-scale objective over the decoder's whole prediction pyramid,
weighted BCE (×30) + IoU (×0.5) + SSIM (×10) per scale, with predictions bilinearly resized to the
label when they do not match. forward also returns the intermediate scales as
intermediate_logits, which is what makes the multi-scale term possible.
Swapping the loss
model.criterion is a plain function attribute, not a submodule, so it never enters the state dict
and you can replace it outright:
from nobg.loss import birefnet_loss, iou_loss, ssim_loss
def my_loss(scaled_preds, gt):
return birefnet_loss(scaled_preds, gt) + 5 * iou_loss(
scaled_preds[-1].sigmoid(), gt
)
model.criterion = my_lossThe signature is (scaled_preds: list[Tensor], gt: Tensor) -> Tensor, where scaled_preds are raw
logits at possibly differing resolutions and gt is (B, 1, H, W) in [0, 1]. See the
losses reference for everything that ships.
SAM3
Start from MultiMatte, and take the processor from the same repo so the tokenizer matches.
Sam3.forward takes labels the same way and defaults to sam3_loss — SAM 3's own
semantic-segmentation objective, weighted focal (×20, alpha=0.6) + dice (×30) — so fine-tuning
matches how the checkpoint was trained. Sam3 always passes a single-element scaled_preds list.
Two differences from BiRefNet:
- Labels are yours to build.
Sam3Processorinheritstransformers' processor, which handles images, text and boxes but notsegmentation_maps. Resize your masks toconfig.image_sizeand pass a(B, 1, H, W)float tensor. - Every batch needs a prompt. Include
input_idsandattention_maskfrom the processor in the collated batch; without themforwardraises.
import torch
import torch.nn.functional as F
from nobg import AutoModel, AutoProcessor
from transformers import Trainer, TrainingArguments
model = AutoModel.from_pretrained("feyninc/multimatte")
processor = AutoProcessor.from_pretrained("feyninc/multimatte")
size = model.config.image_size
def collate(examples):
batch = processor(
images=[example["image"] for example in examples],
text="the main foreground subject",
return_tensors="pt",
)
masks = torch.stack(
[
F.interpolate(
torch.as_tensor(example["mask"], dtype=torch.float32)[None, None]
/ 255,
size=(size, size),
mode="bilinear",
align_corners=False,
)[0]
for example in examples
]
)
return {
"pixel_values": batch["pixel_values"],
"input_ids": batch["input_ids"],
"attention_mask": batch["attention_mask"],
"labels": masks,
}
trainer = Trainer(
model=model,
args=TrainingArguments(
output_dir="outputs",
learning_rate=2e-5,
),
train_dataset=dataset,
data_collator=collate,
)
trainer.train()What sam3_loss deliberately leaves out
SAM 3's full criterion also has Hungarian-matched set losses (box L1/GIoU, classification,
per-instance mask/dice) and a presence-head BCE. Neither applies here: nobg trains on a single
merged matte, so there are no instance-level targets to match, and every training pair has a
foreground, so the presence target is constantly 1 and carries no signal.
Practical notes
- Freeze what you can. These are 0.3–0.84 B models at 1008–1024 px.
peftis in the project'straindependency group for exactly this. - Resolution costs quadratically. Fine-tuning at a smaller
image_sizeand re-parameterizing afterwards is a legitimate strategy. - Evaluate with
nobg.metrics. MAE, S-measure, F/E-measure, boundary IoU and the matting metrics are all pure torch; see Evaluation. - Push the result.
model.push_to_hub(...)andprocessor.push_to_hub(...)write a checkpoint thatAutoModelcan load back; see Hub round-trips.