Encoder: 4 stages of double 3×3 conv blocks (64 → 128 → 256 → 512) with MaxPool2d downsampling, followed by a 1024-channel bottleneck.
Decoder: 4 stages of transposed convolution upsampling with skip concatenation and double conv blocks (512 → 256 → 128 → 64).
Reduction head: extra Conv2d(64→32) + BatchNorm + ReLU before the final 1×1 convolution to reduce capacity and mitigate overfitting.
Output: 1-channel logits (binary segmentation).
All conv blocks use BatchNorm2d + ReLU activation.
Input: 3-channel RGB, 256×256.

Loss Function

CombinedLoss = 0.5 × BCE + 0.5 × Dice Loss

BCE (Binary Cross-Entropy with Logits): provides stable pixel-level gradients and handles class imbalance through its log-based formulation.
Dice Loss: directly optimizes the overlap metric (Dice coefficient), which is the standard evaluation metric for medical segmentation. BCE alone tends to underperform on imbalanced masks where background dominates.
The equal weighting balances pixel-wise accuracy (BCE) with region-level overlap (Dice).

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.7747	—
Test	0.8114	0.2661

Usage

python

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

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

# input: (B, 3, 256, 256) float tensor in [0, 1]
# output: (B, 1, 256, 256) logits — apply sigmoid + threshold for mask

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",
    filename="model.safetensors",
)
model = UNet(in_channels=3, num_classes=1)
model.load_state_dict(load_file(weights))
model.eval()