From 5a0203631845cebdd73ce677c2991d59c71c2f35 Mon Sep 17 00:00:00 2001 From: crosstyan Date: Tue, 10 Mar 2026 00:40:41 +0800 Subject: [PATCH] Add weighted CE proxy and fix loss imports --- ...nt8_sharedalign_weightedce_proxy_1gpu.yaml | 115 ++++++++++++++++++ docs/scoliosis_training_change_log.md | 1 + opengait/modeling/losses/base.py | 20 ++- opengait/modeling/losses/bce.py | 13 +- opengait/modeling/losses/ce.py | 49 +++++++- 5 files changed, 184 insertions(+), 14 deletions(-) create mode 100644 configs/sconet/sconet_scoliosis1k_skeleton_118_sigma15_joint8_sharedalign_weightedce_proxy_1gpu.yaml diff --git a/configs/sconet/sconet_scoliosis1k_skeleton_118_sigma15_joint8_sharedalign_weightedce_proxy_1gpu.yaml b/configs/sconet/sconet_scoliosis1k_skeleton_118_sigma15_joint8_sharedalign_weightedce_proxy_1gpu.yaml new file mode 100644 index 0000000..4991553 --- /dev/null +++ b/configs/sconet/sconet_scoliosis1k_skeleton_118_sigma15_joint8_sharedalign_weightedce_proxy_1gpu.yaml @@ -0,0 +1,115 @@ +data_cfg: + dataset_name: Scoliosis1K + dataset_root: /mnt/public/data/Scoliosis1K/Scoliosis1K-drf-pkl-118-sigma15-joint8-sharedalign + dataset_partition: ./datasets/Scoliosis1K/Scoliosis1K_118.json + data_in_use: + - true + - false + num_workers: 1 + remove_no_gallery: false + test_dataset_name: Scoliosis1K + test_seq_subset_size: 128 + test_seq_subset_seed: 118 + +evaluator_cfg: + enable_float16: true + restore_ckpt_strict: true + restore_hint: 2000 + save_name: ScoNet_skeleton_118_sigma15_joint8_sharedalign_weightedce_proxy_1gpu + eval_func: evaluate_scoliosis + sampler: + batch_shuffle: false + batch_size: 1 + sample_type: all_ordered + type: InferenceSampler + frames_all_limit: 720 + metric: euc + transform: + - type: BaseSilCuttingTransform + +loss_cfg: + - loss_term_weight: 1.0 + margin: 0.2 + type: TripletLoss + log_prefix: triplet + - loss_term_weight: 1.0 + scale: 16 + type: CrossEntropyLoss + log_prefix: softmax + log_accuracy: true + class_weight: + - 1.0 + - 4.0 + - 4.0 + +model_cfg: + model: ScoNet + backbone_cfg: + type: ResNet9 + block: BasicBlock + in_channel: 2 + channels: + - 64 + - 128 + - 256 + - 512 + layers: + - 1 + - 1 + - 1 + - 1 + strides: + - 1 + - 2 + - 2 + - 1 + maxpool: false + SeparateFCs: + in_channels: 512 + out_channels: 256 + parts_num: 16 + SeparateBNNecks: + class_num: 3 + in_channels: 256 + parts_num: 16 + bin_num: + - 16 + +optimizer_cfg: + lr: 0.1 + solver: SGD + weight_decay: 0.0005 + +scheduler_cfg: + gamma: 0.1 + milestones: + - 1000 + - 1500 + - 1800 + scheduler: MultiStepLR + +trainer_cfg: + enable_float16: true + fix_BN: false + with_test: true + log_iter: 100 + restore_ckpt_strict: true + restore_hint: 0 + auto_resume_latest: true + resume_every_iter: 500 + resume_keep: 3 + eval_iter: 500 + save_iter: 2000 + save_name: ScoNet_skeleton_118_sigma15_joint8_sharedalign_weightedce_proxy_1gpu + sync_BN: true + total_iter: 2000 + sampler: + batch_shuffle: true + batch_size: + - 8 + - 8 + frames_num_fixed: 30 + sample_type: fixed_unordered + type: TripletSampler + transform: + - type: BaseSilCuttingTransform diff --git a/docs/scoliosis_training_change_log.md b/docs/scoliosis_training_change_log.md index 62cfa30..c069161 100644 --- a/docs/scoliosis_training_change_log.md +++ b/docs/scoliosis_training_change_log.md @@ -29,6 +29,7 @@ Use it for: | 2026-03-09 | `ScoNet_skeleton_118_sigma15_joint8_sharedalign_nocut_adamw_1gpu_bs8x8` | ScoNet-MT-ske control | `Scoliosis1K-drf-pkl-118-sigma15-joint8-sharedalign` | Switched runtime transform from `BaseSilCuttingTransform` to `BaseSilTransform` (`no-cut`), kept `AdamW`, reduced `8x8` due to 5070 Ti OOM at `12x8` | interrupted | superseded by proxy route before eval | | 2026-03-09 | `ScoNet_skeleton_118_sigma15_joint8_sharedalign_nocut_adamw_proxy_1gpu` | ScoNet-MT-ske proxy | `Scoliosis1K-drf-pkl-118-sigma15-joint8-sharedalign` | Fast proxy route: `no-cut`, `AdamW`, `8x8`, `total_iter=2000`, `eval_iter=500`, `test_seq_subset_size=128` | interrupted | superseded by geometry-fixed proxy before completion | | 2026-03-10 | `ScoNet_skeleton_118_sigma15_joint8_geomfix_proxy_1gpu` | ScoNet-MT-ske proxy | `Scoliosis1K-drf-pkl-118-sigma15-joint8-geomfix` | Geometry ablation: aspect-ratio-preserving crop+pad instead of square-warp resize; `AdamW`, `no-cut`, `8x8`, `total_iter=2000`, `eval_iter=500`, fixed test subset seed `118` | complete | proxy subset unstable: `500 24.22/8.07/33.33/13.00`, `1000 60.16/68.05/58.13/55.25`, `1500 26.56/58.33/35.64/17.68`, `2000 27.34/63.96/37.02/20.14` (Acc/Prec/Rec/F1) | +| 2026-03-10 | `ScoNet_skeleton_118_sigma15_joint8_sharedalign_weightedce_proxy_1gpu` | ScoNet-MT-ske proxy | `Scoliosis1K-drf-pkl-118-sigma15-joint8-sharedalign` | Training-side imbalance ablation: kept the current best shared-align geometry, restored `SGD` baseline settings, and applied weighted CE with class weights `[1.0, 4.0, 4.0]`; `total_iter=2000`, `eval_iter=500`, fixed test subset seed `118` | training | no eval yet | ## Current best skeleton baseline diff --git a/opengait/modeling/losses/base.py b/opengait/modeling/losses/base.py index c235ca0..55f11da 100644 --- a/opengait/modeling/losses/base.py +++ b/opengait/modeling/losses/base.py @@ -1,17 +1,25 @@ +from __future__ import annotations + +from collections.abc import Callable from ctypes import ArgumentError -import torch.nn as nn +from typing import Any + import torch +import torch.nn as nn + from opengait.utils import Odict import functools from opengait.utils import ddp_all_gather -def gather_and_scale_wrapper(func): +def gather_and_scale_wrapper( + func: Callable[..., tuple[torch.Tensor | float, Odict]], +) -> Callable[..., tuple[torch.Tensor | float, Odict]]: """Internal wrapper: gather the input from multple cards to one card, and scale the loss by the number of cards. """ @functools.wraps(func) - def inner(*args, **kwds): + def inner(*args: Any, **kwds: Any) -> tuple[torch.Tensor | float, Odict]: try: for k, v in kwds.items(): @@ -32,7 +40,7 @@ class BaseLoss(nn.Module): Your loss should also subclass this class. """ - def __init__(self, loss_term_weight=1.0): + def __init__(self, loss_term_weight: float = 1.0) -> None: """ Initialize the base class. @@ -43,7 +51,7 @@ class BaseLoss(nn.Module): self.loss_term_weight = loss_term_weight self.info = Odict() - def forward(self, logits, labels): + def forward(self, logits: torch.Tensor, labels: torch.Tensor) -> tuple[torch.Tensor, Odict]: """ The default forward function. @@ -56,4 +64,4 @@ class BaseLoss(nn.Module): Returns: tuple of loss and info. """ - return .0, self.info + return torch.tensor(0.0, device=logits.device), self.info diff --git a/opengait/modeling/losses/bce.py b/opengait/modeling/losses/bce.py index 62034a9..e405e1b 100644 --- a/opengait/modeling/losses/bce.py +++ b/opengait/modeling/losses/bce.py @@ -1,14 +1,21 @@ +from __future__ import annotations + import torch + +from opengait.evaluation.metric import mean_iou +from opengait.utils.common import Odict + from .base import BaseLoss -from evaluation import mean_iou class BinaryCrossEntropyLoss(BaseLoss): - def __init__(self, loss_term_weight=1.0, eps=1.0e-9): + eps: float + + def __init__(self, loss_term_weight: float = 1.0, eps: float = 1.0e-9) -> None: super(BinaryCrossEntropyLoss, self).__init__(loss_term_weight) self.eps = eps - def forward(self, logits, labels): + def forward(self, logits: torch.Tensor, labels: torch.Tensor) -> tuple[torch.Tensor, Odict]: """ logits: [n, 1, h, w] labels: [n, 1, h, w] diff --git a/opengait/modeling/losses/ce.py b/opengait/modeling/losses/ce.py index 8740fd0..fce602b 100644 --- a/opengait/modeling/losses/ce.py +++ b/opengait/modeling/losses/ce.py @@ -1,29 +1,68 @@ +from __future__ import annotations + +from collections.abc import Sequence + +import torch import torch.nn.functional as F +from opengait.utils.common import Odict + from .base import BaseLoss class CrossEntropyLoss(BaseLoss): - def __init__(self, scale=2**4, label_smooth=True, eps=0.1, loss_term_weight=1.0, log_accuracy=False): + scale: float + label_smooth: bool + eps: float + log_accuracy: bool + class_weight: torch.Tensor | None = None + + def __init__( + self, + scale: float = 2**4, + label_smooth: bool = True, + eps: float = 0.1, + loss_term_weight: float = 1.0, + log_accuracy: bool = False, + class_weight: Sequence[float] | None = None, + ) -> None: super(CrossEntropyLoss, self).__init__(loss_term_weight) self.scale = scale self.label_smooth = label_smooth self.eps = eps self.log_accuracy = log_accuracy + weight_tensor = ( + None + if class_weight is None + else torch.as_tensor(class_weight, dtype=torch.float32) + ) + if class_weight is None: + self.register_buffer("class_weight", weight_tensor) + else: + self.register_buffer("class_weight", weight_tensor) - def forward(self, logits, labels): + def forward(self, logits: torch.Tensor, labels: torch.Tensor) -> tuple[torch.Tensor, Odict]: """ logits: [n, c, p] labels: [n] """ - n, c, p = logits.size() + _n, _c, p = logits.size() logits = logits.float() labels = labels.unsqueeze(1) + class_weight = self.class_weight if isinstance(self.class_weight, torch.Tensor) else None if self.label_smooth: loss = F.cross_entropy( - logits*self.scale, labels.repeat(1, p), label_smoothing=self.eps) + logits * self.scale, + labels.repeat(1, p), + weight=class_weight, + label_smoothing=self.eps, + ) else: - loss = F.cross_entropy(logits*self.scale, labels.repeat(1, p)) + loss = F.cross_entropy( + logits * self.scale, + labels.repeat(1, p), + weight=class_weight, + ) self.info.update({'loss': loss.detach().clone()}) if self.log_accuracy: pred = logits.argmax(dim=1) # [n, p]