Vagneralexx/ben2-ft-mix

BEN2 Fine-Tuned (cons4) - Background Removal

Fine-tuned BEN2 Base (MIT license) for high-quality background removal.

Variants

Variant	File	Precision	Recommended Use
GPU	gpu_fp16/pytorch_model.bin	fp16	NVIDIA GPU (RTX 20/30/40 series)
CPU	cpu_fp32/pytorch_model.bin	fp32	CPU (x86_64, ARM64)

Performance

Dataset	Mean IoU	Median	P10	<0.85
Photos (200)	0.9960	0.9978	0.9918	0/200
Arts (122)	0.9175	0.9923	0.7856	14/122

Quick Start

GPU (fp16)

import torch
from model import BEN_Base

model = BEN_Base()
model.load_state_dict(torch.load("gpu_fp16/pytorch_model.bin", map_location="cuda"))
model.to("cuda").eval().half()

# Inference
with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16):
    logits = model(input_tensor)
    mask = torch.sigmoid(logits).sigmoid()

CPU (fp32)

import torch
from model import BEN_Base

model = BEN_Base()
model.load_state_dict(torch.load("cpu_fp32/pytorch_model.bin", map_location="cpu"))
model.to("cpu").eval()

# torch.compile for CPU speedup (PyTorch 2.0+)
model = torch.compile(model)

with torch.no_grad():
    logits = model(input_tensor)
    mask = torch.sigmoid(logits)

Training Details

• Base: BEN2 Base (MIT) from D:/tmp/ben2-treino/BEN2_Base.pth
• Fine-tuning: 4 epochs (cons4), mix of 9.4k photos (P3M) + 122 arts + 107 extras
• Loss: BCE + IoU + edge-weighted BCE
• Optimizer: AdamW, lr=3e-5 (backbone 5e-6), accum=4, batch=1
• Hardware: NVIDIA RTX (WebGPU training), fp16

License

MIT (inherited from BEN2 Base)