Add weighted CE proxy and fix loss imports

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

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
 

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

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]

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]