CrossEntropy: replaces BCE for the multi-class case, providing per-pixel classification gradients across all classes.
Multi-class Dice Loss: computes per-class Dice using softmax probabilities and one-hot targets, then averages across classes. Ensures balanced optimization even with class imbalance.

Training Setup

Parameter	Value
Epochs	20
Optimizer	Adam
Learning rate	1e-4
LR scheduler	ReduceLROnPlateau (factor=0.5, patience=3, min_lr=1e-6)
Batch size	8
Image size	256×256
GPU	NVIDIA GeForce RTX 3080 Ti
Dataset splits	train=800, val=100, test=100

Metrics

Split	Dice Coefficient	Loss
Validation (best)	0.9036	—
Test	0.9169	0.3322

Usage

python

import torch
from model import UNet
from safetensors.torch import load_file

model = UNet(in_channels=3, num_classes=3)
model.load_state_dict(load_file("model.safetensors"))
model.eval()

# input: (B, 3, 256, 256) float tensor in [0, 1]
# output: (B, 3, 256, 256) logits — apply argmax(dim=1) for class map

Or download from Hub:

python

from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from model import UNet

weights = hf_hub_download(
    repo_id="sebastiao-teixeira/week04-polyp-segmentation-unet-multiclass",
    filename="model.safetensors",
)
model = UNet(in_channels=3, num_classes=3)
model.load_state_dict(load_file(weights))
model.eval()