530 lines
20 KiB
Python
530 lines
20 KiB
Python
#python _8_train_segformer_b3.py --epochs 180 --batch 2 --lr 3e-5 --wd 0.01 --num_workers 4 --amp --amp_val --grad_accum 2 --class_weights auto --main_class navegavel --resume
|
|
|
|
# _8_train_segformer_b3.py (PATCH)
|
|
import os
|
|
|
|
# (Opcional) ajuda com fragmentação em algumas máquinas.
|
|
# Idealmente isso deveria vir ANTES de importar torch, mas já ajuda quando setado fora também.
|
|
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:128")
|
|
|
|
import json
|
|
import math
|
|
import time
|
|
import argparse
|
|
from typing import Dict, List, Optional, Tuple, Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.data import DataLoader
|
|
|
|
# >>> troque AMP deprecated
|
|
from torch.amp import autocast, GradScaler
|
|
|
|
from transformers import SegformerForSemanticSegmentation
|
|
|
|
# >>>>>>>>> AJUSTE AQUI <<<<<<<<<
|
|
from _8_train_fastscnn_v2 import ROISegDataset # troque se necessário
|
|
|
|
|
|
def set_seed(seed: int = 42):
|
|
import random
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
torch.cuda.manual_seed_all(seed)
|
|
|
|
|
|
def _ids_by_names(class_names: Any, wanted_names: List[str]) -> List[int]:
|
|
if class_names is None:
|
|
return []
|
|
wanted_names = [str(x).lower() for x in wanted_names]
|
|
|
|
if isinstance(class_names, list):
|
|
low = [c.lower() for c in class_names]
|
|
return [low.index(n) for n in wanted_names if n in low]
|
|
|
|
if isinstance(class_names, dict):
|
|
if all(isinstance(k, int) for k in class_names.keys()):
|
|
inv = {str(v).lower(): int(k) for k, v in class_names.items()}
|
|
return [inv[n] for n in wanted_names if n in inv]
|
|
inv = {str(k).lower(): int(v) for k, v in class_names.items()}
|
|
return [inv[n] for n in wanted_names if n in inv]
|
|
|
|
return []
|
|
|
|
|
|
@torch.no_grad()
|
|
def update_confusion_matrix(cm: torch.Tensor, preds: torch.Tensor, labels: torch.Tensor, num_classes: int, ignore_index: int = 255):
|
|
preds = preds.view(-1)
|
|
labels = labels.view(-1)
|
|
|
|
valid = labels != ignore_index
|
|
preds = preds[valid]
|
|
labels = labels[valid]
|
|
|
|
idx = labels * num_classes + preds
|
|
bins = torch.bincount(idx, minlength=num_classes * num_classes)
|
|
cm += bins.view(num_classes, num_classes)
|
|
|
|
|
|
@torch.no_grad()
|
|
def compute_iou_from_cm(cm: torch.Tensor, eps: float = 1e-6) -> Tuple[float, List[float]]:
|
|
cm = cm.float()
|
|
tp = torch.diag(cm)
|
|
fp = cm.sum(0) - tp
|
|
fn = cm.sum(1) - tp
|
|
denom = tp + fp + fn + eps
|
|
iou = (tp / denom).cpu().tolist()
|
|
miou = float(np.mean(iou))
|
|
return miou, iou
|
|
|
|
|
|
@torch.no_grad()
|
|
def compute_pixel_acc_from_cm(cm: torch.Tensor, eps: float = 1e-6) -> float:
|
|
cm = cm.float()
|
|
acc = (torch.diag(cm).sum() / (cm.sum() + eps)).item()
|
|
return acc
|
|
|
|
|
|
IMAGENET_MEAN = [0.485, 0.456, 0.406]
|
|
IMAGENET_STD = [0.229, 0.224, 0.225]
|
|
NORM_MEAN = torch.tensor(IMAGENET_MEAN).view(3, 1, 1)
|
|
NORM_STD = torch.tensor(IMAGENET_STD).view(3, 1, 1).clamp_min(1e-6)
|
|
def normalize_img(img: torch.Tensor) -> torch.Tensor:
|
|
return (img - NORM_MEAN.to(img.device)) / NORM_STD.to(img.device)
|
|
|
|
|
|
def default_collate(batch):
|
|
imgs = []
|
|
masks = []
|
|
for item in batch:
|
|
if isinstance(item, dict):
|
|
img = item["image"]
|
|
mask = item["mask"]
|
|
else:
|
|
img, mask = item
|
|
|
|
if isinstance(img, np.ndarray):
|
|
img = torch.from_numpy(img)
|
|
if isinstance(mask, np.ndarray):
|
|
mask = torch.from_numpy(mask)
|
|
|
|
if img.ndim == 3 and img.shape[-1] == 3:
|
|
img = img.permute(2, 0, 1)
|
|
|
|
if img.dtype != torch.float32:
|
|
img = img.float()
|
|
if img.max() > 1.5:
|
|
img = img / 255.0
|
|
|
|
imgs.append(img)
|
|
masks.append(mask.long())
|
|
|
|
return torch.stack(imgs, 0), torch.stack(masks, 0)
|
|
|
|
|
|
def estimate_class_weights(ds: ROISegDataset, num_classes: int, ignore_index: int = 255, max_samples: int = 800) -> torch.Tensor:
|
|
n = min(len(ds), max_samples)
|
|
idxs = np.random.choice(len(ds), size=n, replace=False)
|
|
|
|
counts = np.zeros(num_classes, dtype=np.float64)
|
|
for i in idxs:
|
|
item = ds[i]
|
|
mask = item["mask"] if isinstance(item, dict) else item[1]
|
|
m = mask.cpu().numpy() if isinstance(mask, torch.Tensor) else np.array(mask)
|
|
|
|
m = m.reshape(-1)
|
|
m = m[m != ignore_index]
|
|
if m.size == 0:
|
|
continue
|
|
|
|
counts += np.bincount(m, minlength=num_classes)[:num_classes]
|
|
|
|
freq = counts / (counts.sum() + 1e-12)
|
|
freq = np.clip(freq, 1e-12, 1.0)
|
|
weights = 1.0 / np.log(1.02 + freq)
|
|
weights = weights / weights.mean()
|
|
return torch.tensor(weights, dtype=torch.float32)
|
|
|
|
|
|
def save_checkpoint(path: str,
|
|
model: nn.Module,
|
|
optimizer: torch.optim.Optimizer,
|
|
scaler: Optional[GradScaler],
|
|
epoch: int,
|
|
best_miou: float,
|
|
best_main_iou: float,
|
|
extra: Optional[Dict[str, Any]] = None):
|
|
ckpt = {
|
|
"epoch": epoch,
|
|
"model": model.state_dict(),
|
|
"optimizer": optimizer.state_dict(),
|
|
"best_miou": best_miou,
|
|
"best_main_iou": best_main_iou,
|
|
}
|
|
if scaler is not None:
|
|
ckpt["scaler"] = scaler.state_dict()
|
|
if extra:
|
|
ckpt["extra"] = extra
|
|
torch.save(ckpt, path)
|
|
|
|
|
|
def load_checkpoint(path: str,
|
|
model: nn.Module,
|
|
optimizer: Optional[torch.optim.Optimizer] = None,
|
|
scaler: Optional[GradScaler] = None,
|
|
map_location: str = "cpu") -> Dict[str, Any]:
|
|
ckpt = torch.load(path, map_location=map_location, weights_only=False)
|
|
model.load_state_dict(ckpt["model"], strict=True)
|
|
if optimizer is not None and "optimizer" in ckpt:
|
|
optimizer.load_state_dict(ckpt["optimizer"])
|
|
if scaler is not None and "scaler" in ckpt:
|
|
scaler.load_state_dict(ckpt["scaler"])
|
|
return ckpt
|
|
|
|
|
|
def run_one_epoch(model: nn.Module,
|
|
loader: DataLoader,
|
|
optimizer: Optional[torch.optim.Optimizer],
|
|
device: torch.device,
|
|
num_classes: int,
|
|
ignore_index: int,
|
|
criterion: nn.Module,
|
|
amp: bool,
|
|
scaler: Optional[GradScaler],
|
|
train: bool,
|
|
grad_accum: int = 1) -> Dict[str, Any]:
|
|
|
|
model.train(train)
|
|
|
|
total_loss = 0.0
|
|
cm = torch.zeros((num_classes, num_classes), dtype=torch.int64, device=device)
|
|
|
|
t0 = time.time()
|
|
n_batches = 0
|
|
|
|
# >>>>> ESSENCIAL: desliga grad no val
|
|
with torch.set_grad_enabled(train):
|
|
if train and optimizer is not None:
|
|
optimizer.zero_grad(set_to_none=True)
|
|
|
|
for step, (imgs, masks) in enumerate(loader):
|
|
imgs = imgs.to(device, non_blocking=True)
|
|
masks = masks.to(device, non_blocking=True)
|
|
|
|
imgs = normalize_img(imgs)
|
|
|
|
# AMP tanto em train quanto (opcionalmente) em val
|
|
with autocast(device_type="cuda", enabled=amp and device.type == "cuda"):
|
|
out = model(pixel_values=imgs)
|
|
logits = out.logits
|
|
|
|
if logits.shape[-2:] != masks.shape[-2:]:
|
|
logits = torch.nn.functional.interpolate(
|
|
logits, size=masks.shape[-2:], mode="bilinear", align_corners=False
|
|
)
|
|
|
|
loss = criterion(logits, masks)
|
|
if train and grad_accum > 1:
|
|
loss = loss / grad_accum
|
|
|
|
if train and optimizer is not None:
|
|
if amp and scaler is not None and device.type == "cuda":
|
|
scaler.scale(loss).backward()
|
|
if ((step + 1) % grad_accum) == 0:
|
|
scaler.step(optimizer)
|
|
scaler.update()
|
|
optimizer.zero_grad(set_to_none=True)
|
|
else:
|
|
loss.backward()
|
|
if ((step + 1) % grad_accum) == 0:
|
|
optimizer.step()
|
|
optimizer.zero_grad(set_to_none=True)
|
|
|
|
total_loss += float(loss.item()) * (grad_accum if train and grad_accum > 1 else 1.0)
|
|
n_batches += 1
|
|
|
|
preds = torch.argmax(logits, dim=1)
|
|
update_confusion_matrix(cm, preds, masks, num_classes=num_classes, ignore_index=ignore_index)
|
|
|
|
dt = time.time() - t0
|
|
avg_loss = total_loss / max(n_batches, 1)
|
|
miou, iou_per_class = compute_iou_from_cm(cm)
|
|
acc = compute_pixel_acc_from_cm(cm)
|
|
|
|
return {"loss": avg_loss, "miou": miou, "iou_per_class": iou_per_class, "acc": acc, "time_s": dt}
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", default="config.json")
|
|
parser.add_argument("--epochs", type=int, default=120)
|
|
parser.add_argument("--batch", type=int, default=1) # <<< default seguro pra 8GB
|
|
parser.add_argument("--lr", type=float, default=6e-5)
|
|
parser.add_argument("--wd", type=float, default=0.01)
|
|
parser.add_argument("--num_workers", type=int, default=4)
|
|
|
|
parser.add_argument("--save_every", type=int, default=10)
|
|
parser.add_argument("--resume", action="store_true")
|
|
|
|
parser.add_argument("--amp", action="store_true")
|
|
parser.add_argument("--amp_val", action="store_true") # <<< novo: AMP no val
|
|
parser.add_argument("--grad_accum", type=int, default=1) # <<< novo
|
|
|
|
parser.add_argument("--grad_ckpt", action="store_true") # <<< novo: checkpointing
|
|
parser.add_argument("--ignore_index", type=int, default=255)
|
|
parser.add_argument("--class_weights", type=str, default="auto")
|
|
parser.add_argument("--main_class", type=str, default=None)
|
|
parser.add_argument("--es_classes", type=str, default="")
|
|
parser.add_argument("--seed", type=int, default=42)
|
|
parser.add_argument("--norm_stats", type=str, default=None, help="Caminho para JSON com mean/std por canal (ex: norm_stats.json).")
|
|
args = parser.parse_args()
|
|
|
|
set_seed(args.seed)
|
|
|
|
with open(args.config, "r") as f:
|
|
config = json.load(f)
|
|
|
|
MODELO = config["camera"]
|
|
MODEL_NAME = config["model_name"]
|
|
RESOLUCAO = config["resolucao"]
|
|
ROI_INICIO = config["roi_inicio"]
|
|
ROI_TAMANHO = config["roi_tamanho"]
|
|
MAIN_CLASS_NAME = str(config.get("main_class_name", "erva")).lower()
|
|
BACKBONE = config["backbone"]
|
|
if args.main_class is not None:
|
|
MAIN_CLASS_NAME = args.main_class.lower()
|
|
|
|
dataset_path = os.path.join(MODELO, "dataset")
|
|
labelmap_path = os.path.join(dataset_path, "labelmap.txt")
|
|
|
|
save_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME)
|
|
os.makedirs(save_path, exist_ok=True)
|
|
last_ckpt_path = os.path.join(save_path, "last.pt")
|
|
best_miou_path = os.path.join(save_path, "best_miou.pt")
|
|
best_main_path = os.path.join(save_path, "best_main.pt")
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
print(f"Device: {device}")
|
|
|
|
# Datasets
|
|
ds_train = ROISegDataset(
|
|
os.path.join(dataset_path, "split", "train"),
|
|
save_path, ROI_INICIO, ROI_TAMANHO,
|
|
RESOLUCAO[0], RESOLUCAO[1], labelmap_path
|
|
)
|
|
ds_val = ROISegDataset(
|
|
os.path.join(dataset_path, "split", "val"),
|
|
save_path, ROI_INICIO, ROI_TAMANHO,
|
|
RESOLUCAO[0], RESOLUCAO[1], labelmap_path
|
|
)
|
|
|
|
CLASS_NAMES = getattr(ds_train, "classes", None)
|
|
if CLASS_NAMES is None:
|
|
raise RuntimeError("ROISegDataset precisa expor .classes (list ou dict).")
|
|
|
|
if isinstance(CLASS_NAMES, list):
|
|
num_classes = len(CLASS_NAMES)
|
|
class_id_by_name = {n.lower(): i for i, n in enumerate(CLASS_NAMES)}
|
|
class_name_by_id = {i: n for i, n in enumerate(CLASS_NAMES)}
|
|
elif isinstance(CLASS_NAMES, dict):
|
|
if all(isinstance(k, int) for k in CLASS_NAMES.keys()):
|
|
num_classes = len(CLASS_NAMES)
|
|
class_name_by_id = {int(k): str(v) for k, v in CLASS_NAMES.items()}
|
|
class_id_by_name = {str(v).lower(): int(k) for k, v in CLASS_NAMES.items()}
|
|
else:
|
|
class_id_by_name = {str(k).lower(): int(v) for k, v in CLASS_NAMES.items()}
|
|
num_classes = len(class_id_by_name)
|
|
class_name_by_id = {v: k for k, v in class_id_by_name.items()}
|
|
else:
|
|
raise RuntimeError("Formato de .classes não reconhecido.")
|
|
|
|
main_class_id = class_id_by_name.get(MAIN_CLASS_NAME, None)
|
|
if main_class_id is None:
|
|
print(f"[WARN] main_class_name='{MAIN_CLASS_NAME}' não encontrado. best_main_iou usa mIoU.")
|
|
else:
|
|
print(f"Main class: '{MAIN_CLASS_NAME}' -> id={main_class_id}")
|
|
|
|
es_names = [x.strip().lower() for x in args.es_classes.split(",") if x.strip()]
|
|
ES_CLASS_IDS = _ids_by_names(CLASS_NAMES, es_names)
|
|
|
|
dl_train = DataLoader(
|
|
ds_train, batch_size=args.batch, shuffle=True,
|
|
num_workers=args.num_workers, pin_memory=True,
|
|
collate_fn=default_collate, drop_last=True
|
|
)
|
|
dl_val = DataLoader(
|
|
ds_val, batch_size=1, shuffle=False, # <<< val com batch 1 é mais estável
|
|
num_workers=max(2, args.num_workers // 2), pin_memory=True,
|
|
collate_fn=default_collate, drop_last=False
|
|
)
|
|
|
|
model = SegformerForSemanticSegmentation.from_pretrained(
|
|
BACKBONE,
|
|
num_labels=num_classes,
|
|
ignore_mismatched_sizes=True,
|
|
use_safetensors=True
|
|
)
|
|
|
|
if args.grad_ckpt:
|
|
try:
|
|
# Nem toda versão/classe suporta
|
|
model.gradient_checkpointing_enable()
|
|
print("[OK] gradient checkpointing enabled")
|
|
except Exception as e:
|
|
print(f"[WARN] gradient checkpointing não suportado aqui ({type(e).__name__}: {e}). Seguindo sem.")
|
|
|
|
model.to(device)
|
|
|
|
# Loss weights
|
|
if args.class_weights.lower() == "none":
|
|
weights = None
|
|
elif args.class_weights.lower() == "auto":
|
|
w = estimate_class_weights(ds_train, num_classes=num_classes, ignore_index=args.ignore_index)
|
|
weights = w.to(device)
|
|
print("Class weights (auto):", w.cpu().numpy().round(3).tolist())
|
|
else:
|
|
parts = [float(x) for x in args.class_weights.split(",")]
|
|
if len(parts) != num_classes:
|
|
raise ValueError(f"class_weights manual precisa ter {num_classes} valores, recebeu {len(parts)}.")
|
|
weights = torch.tensor(parts, dtype=torch.float32, device=device)
|
|
|
|
criterion = nn.CrossEntropyLoss(weight=weights, ignore_index=args.ignore_index)
|
|
|
|
optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.wd)
|
|
|
|
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
|
|
optimizer, mode="min", factor=0.5, patience=6, threshold=1e-4, verbose=True
|
|
)
|
|
|
|
scaler = GradScaler(enabled=args.amp and device.type == "cuda")
|
|
|
|
|
|
|
|
# ==========================
|
|
# Normalizador (fixo ou dinâmico)
|
|
# ==========================
|
|
global NORM_MEAN, NORM_STD
|
|
norm_stats = None
|
|
|
|
# Caminho padrão: dentro do dataset, nome do arquivo de stats
|
|
# (ajusta aqui pro nome que você realmente usou: norm_stats.json, por ex.)
|
|
norm_stats_path = os.path.join(save_path, "norm_stats.json")
|
|
if args.norm_stats is not None:
|
|
norm_stats_path = args.norm_stats
|
|
|
|
if norm_stats_path is not None and os.path.exists(norm_stats_path):
|
|
with open(norm_stats_path, "r", encoding="utf-8") as f:
|
|
norm_stats = json.load(f)
|
|
|
|
stats_channels = norm_stats.get("channels", [])
|
|
stats_mean = norm_stats.get("mean", [])
|
|
stats_std = norm_stats.get("std", [])
|
|
|
|
print(f"[NORM] usando stats fixos de: {norm_stats_path}")
|
|
print(f"[NORM] channels={stats_channels}")
|
|
print(f"[NORM] mean={stats_mean}")
|
|
print(f"[NORM] std ={stats_std}")
|
|
|
|
# Garante que temos pelo menos R,G,B
|
|
idx_by_name = {name: i for i, name in enumerate(stats_channels)}
|
|
required = ["R", "G", "B"]
|
|
if not all(ch in idx_by_name for ch in required):
|
|
print("[NORM] AVISO: norm_stats não contém todos os canais R,G,B. Mantendo normalize imagenet.")
|
|
else:
|
|
NORM_MEAN = torch.tensor(stats_mean, dtype=torch.float32, device=device).view(3, 1, 1)
|
|
NORM_STD = torch.tensor(stats_std, dtype=torch.float32, device=device).view(3, 1, 1).clamp_min(1e-6)
|
|
|
|
print("[NORM] Normalização fixa por canal ativada para [R,G,B].")
|
|
else:
|
|
if norm_stats_path:
|
|
print(f"[NORM] Caminho de norm_stats não encontrado: {norm_stats_path}. Usando normalize imagenet.")
|
|
else:
|
|
print("[NORM] norm_stats não informado. Usando normalize imagenet.")
|
|
|
|
|
|
|
|
start_epoch = 1
|
|
best_miou = -1.0
|
|
best_main_iou = -1.0
|
|
|
|
if args.resume and os.path.exists(best_miou_path):
|
|
ckpt = load_checkpoint(best_miou_path, model, optimizer, scaler=scaler, map_location="cpu")
|
|
start_epoch = int(ckpt["epoch"]) + 1
|
|
best_miou = float(ckpt.get("best_miou", -1.0))
|
|
best_main_iou = float(ckpt.get("best_main_iou", -1.0))
|
|
print(f"[RESUME] epoch={start_epoch} best_miou={best_miou:.4f} best_main_iou={best_main_iou:.4f}")
|
|
|
|
def pretty_iou(iou_list):
|
|
return " | ".join([f"{class_name_by_id.get(cid,cid)}:{v:.3f}" for cid, v in enumerate(iou_list)])
|
|
|
|
for epoch in range(start_epoch, args.epochs + 1):
|
|
lr_now = optimizer.param_groups[0]["lr"]
|
|
print(f"\n==== Epoch {epoch}/{args.epochs} | lr={lr_now:.2e} ====")
|
|
|
|
if device.type == "cuda":
|
|
torch.cuda.empty_cache()
|
|
|
|
tr = run_one_epoch(
|
|
model=model, loader=dl_train, optimizer=optimizer,
|
|
device=device, num_classes=num_classes, ignore_index=args.ignore_index,
|
|
criterion=criterion, amp=args.amp, scaler=scaler, train=True,
|
|
grad_accum=max(1, args.grad_accum)
|
|
)
|
|
|
|
if device.type == "cuda":
|
|
torch.cuda.empty_cache()
|
|
|
|
va = run_one_epoch(
|
|
model=model, loader=dl_val, optimizer=None,
|
|
device=device, num_classes=num_classes, ignore_index=args.ignore_index,
|
|
criterion=criterion, amp=args.amp_val, scaler=None, train=False,
|
|
grad_accum=1
|
|
)
|
|
|
|
scheduler.step(va["loss"])
|
|
|
|
main_iou = va["miou"] if main_class_id is None else va["iou_per_class"][main_class_id]
|
|
|
|
print(f"TRAIN: loss={tr['loss']:.4f} acc={tr['acc']:.4f} miou={tr['miou']:.4f} (t={tr['time_s']:.1f}s)")
|
|
print(f"VAL : loss={va['loss']:.4f} acc={va['acc']:.4f} miou={va['miou']:.4f} main_iou={main_iou:.4f} (t={va['time_s']:.1f}s)")
|
|
print("IoU per class:", pretty_iou(va["iou_per_class"]))
|
|
|
|
if ES_CLASS_IDS:
|
|
mini = " | ".join([f"{class_name_by_id[cid]}:{va['iou_per_class'][cid]:.3f}" for cid in ES_CLASS_IDS])
|
|
print("ES classes:", mini)
|
|
|
|
save_checkpoint(
|
|
last_ckpt_path, model, optimizer, scaler=scaler,
|
|
epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou,
|
|
extra={"val_loss": va["loss"], "val_miou": va["miou"], "val_main_iou": float(main_iou)}
|
|
)
|
|
|
|
if args.save_every > 0 and (epoch % args.save_every == 0):
|
|
save_checkpoint(
|
|
os.path.join(save_path, f"epoch_{epoch:04d}.pt"),
|
|
model, optimizer, scaler=scaler,
|
|
epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou
|
|
)
|
|
|
|
if va["miou"] > best_miou:
|
|
best_miou = va["miou"]
|
|
save_checkpoint(
|
|
best_miou_path, model, optimizer, scaler=scaler,
|
|
epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou
|
|
)
|
|
print(f"[BEST mIoU] {best_miou:.4f} -> saved: {best_miou_path}")
|
|
|
|
if float(main_iou) > best_main_iou:
|
|
best_main_iou = float(main_iou)
|
|
save_checkpoint(
|
|
best_main_path, model, optimizer, scaler=scaler,
|
|
epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou
|
|
)
|
|
print(f"[BEST MAIN] {best_main_iou:.4f} -> saved: {best_main_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|