File size: 24,307 Bytes
e8f664f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 | # Phân tích source code: các điểm có thể làm giảm hiệu năng và hướng tối ưu
Tài liệu này tổng hợp kết quả rà soát source code trong `code/` của repo, tập trung vào hai nghĩa của "hiệu năng":
- **Hiệu năng mô hình**: Dice/JC/HD/ASD, chất lượng pseudo-label, khả năng tổng quát qua domain shift.
- **Hiệu năng mã nguồn/vận hành**: tốc độ train/eval, tính ổn định, khả năng reproduce, khả năng mở rộng dự án.
Kết luận ngắn: **có thể cải thiện dự án đáng kể bằng cách tối ưu mã nguồn**. Những việc nên ưu tiên không phải refactor lớn, mà là sửa một số điểm trong loss, evaluation, EMA teacher, split labeled data, scheduler và SAM/LoRA loading.
## Phạm vi đã đọc
Đã rà soát các file Python chính trong:
- `code/train.py`, `code/test.py`
- `code/dataloaders/`
- `code/utils/`
- `code/networks/`
- `code/segment_anything/`
- `code/sam_lora_image_encoder.py`
- `code/summarize_logs.py` và test liên quan
Không xem các file `.pyc`, dữ liệu mẫu trong `data_format/`, và các bản snapshot train trong `model/*/train/*/train.py` như source chính, vì chúng là output/thực nghiệm được sinh ra từ script.
## Tổng kết ưu tiên
| Ưu tiên | Khu vực | Vấn đề | Ảnh hưởng | Hướng xử lý |
| --- | --- | --- | --- | --- |
| P0 | Loss semi-supervised | Mask Dice unsupervised xử lý sai class nền | Pseudo-label low-confidence vẫn tác động vào background, có thể kéo Dice xuống | Sửa `DiceLossWithMask._one_hot_mask_encoder()` để confidence mask được áp dụng cho mọi class |
| P0 | Loss semi-supervised | CE unsupervised chia trung bình trên toàn ảnh | Loss thay đổi theo mask ratio, threshold cao làm tín hiệu unsup gần biến mất | Chuẩn hóa theo số pixel hợp lệ: `sum(masked_loss) / mask.sum().clamp_min(1)` |
| P0 | Evaluation | `test.py` load sai đường dẫn checkpoint so với `train.py` | Dễ đánh giá nhầm hoặc fail checkpoint | Chuẩn hóa `experiment_dir`, hỗ trợ `train/<save_name>` hoặc tự động tìm checkpoint |
| P0 | Metric | HD/ASD không xử lý đầy đủ mask rỗng | Có thể crash hoặc phạt điểm sai cho slice không có object | Xử lý 4 case: both empty, pred empty, gt empty, both non-empty |
| P1 | Training | EMA teacher đang ở `train()` khi sinh pseudo-label | BN/running stats của teacher bị nhiễm batch unlabeled, teacher kém ổn định | Đặt EMA teacher `eval()` trong bước pseudo-label |
| P1 | LR schedule | Warmup hiện tại gần như không warmup thật | SAM có thể học quá chậm hoặc quá gắt tùy flag | Thêm scheduler theo iteration cho SAM và U-Net |
| P1 | Data split | Labeled set lấy `first N` sau sort file | Labeled data có thể lệch patient/slice/domain | Shuffle/stratified split bằng seed và lưu split file |
| P1 | Metric Dice | Smoothing Dice `+1` làm inflate object nhỏ/false positive | Metric có thể lạc quan giả, nhất là BUSI/cup nhỏ | Dùng epsilon nhỏ và policy both-empty rõ ràng |
| P2 | SAM/LoRA | SAM train không prompt, `multimask_output` bị bỏ qua | Mất lợi thế promptable segmentation của SAM | Thử prompt từ label/pseudo-label hoặc đặt tên rõ "semantic mode" |
| P2 | Runtime | CPU/NumPy/Scipy trong training loop | Train chậm, GPU bị idle | Đưa confidence/resize/metric sang torch GPU nếu có thể |
| P2 | Data transform | Resize image cho U-Net dùng nearest-neighbor | Ảnh đầu vào bị vỡ hạt/aliasing, có thể giảm Dice U-Net | Dùng bilinear/area cho image, nearest chỉ cho label |
| P3 | Maintainability | Một số import/debug dependency, API checkpoint, DSBN | Khó deploy/mở rộng | Dọn import, checkpoint dict đầy đủ, assert domain batch |
## Phát hiện chi tiết
### 1. Mask Dice unsupervised đang sai với confidence mask
Trong [code/utils/losses.py](../code/utils/losses.py), `DiceLossWithMask._one_hot_mask_encoder()` hiện tại:
```python
temp_prob = input_tensor * i == i * torch.ones_like(input_tensor)
```
Khi `i = 0`, điều kiện này luôn đúng, vì `input_tensor * 0 == 0`. Nghĩa là class background luôn được tính Dice trên toàn bộ ảnh, kể cả vùng low-confidence đã bị threshold loại bỏ.
Nó được dùng trong unsupervised loss tại [code/train.py](../code/train.py#L621) và [code/train.py](../code/train.py#L623). Nếu confidence mask có ý nghĩa là "chỉ học từ pixel đủ tin cậy", cách hiện tại phá vỡ ý nghĩa đó cho class nền.
Khuyến nghị:
- Nếu `mask` là confidence mask dạng shape `B,1,H,W`, hãy expand/repeat mask cho tất cả class.
- Nếu cần mask riêng theo class, nên tạo mask từ `target_onehot * confidence_mask`, không so sánh `input_tensor * i`.
- Thêm unit test nhỏ cho `DiceLossWithMask`: với mask toàn 0 thì Dice unsup phải gần 0 contribution cho mọi class.
Mức độ tác động: **cao**. Đây là điểm có khả năng ảnh hưởng trực tiếp đến Dice vì nằm ngay trong tín hiệu học từ unlabeled data.
### 2. CE unsupervised bị scale theo mask ratio
Trong [code/train.py](../code/train.py#L621) và [code/train.py](../code/train.py#L623):
```python
(ce_loss(logits, pseudo_label) * mask.squeeze(1)).mean()
```
`mean()` chia theo toàn bộ pixel, không chia theo số pixel được chọn. Khi threshold cao, mask nhỏ, loss bị nhỏ đi rất mạnh. Khi mask ratio dao động theo iteration/domain, scale của unsupervised objective cũng dao động theo, làm training khó ổn định.
Khuyến nghị:
```python
pixel_loss = ce_loss(logits, pseudo_label)
valid = mask.squeeze(1)
masked_ce = (pixel_loss * valid).sum() / valid.sum().clamp_min(1.0)
```
Có thể log riêng `mask_ratio` và `masked_ce` để xem tín hiệu unsup có thực sự đang học hay không.
Mức độ tác động: **cao**.
### 3. EMA teacher đang để `train()` khi tạo pseudo-label
Trong [code/train.py](../code/train.py#L471)-[code/train.py](../code/train.py#L474), cả student và EMA model đều được đặt `train()` mỗi epoch. Sau đó pseudo-label được tạo trong `torch.no_grad()` tại [code/train.py](../code/train.py#L527). `no_grad()` không ngăn BatchNorm cập nhật running mean/var nếu model ở train mode.
Với U-Net có BatchNorm trong [code/networks/unet_parts.py](../code/networks/unet_parts.py#L15), EMA teacher có thể bị "nhiễm" thống kê batch unlabeled, làm teacher kém ổn định.
Khuyến nghị:
- Student: `SAM_model.train()`, `unet_model.train()`.
- Teacher EMA khi sinh pseudo-label: `ema_SAM_model.eval()`, `ema_unet_model.eval()`.
- Nếu vẫn cần train mode cho một module đặc biệt, tách rõ module đó và test ablation.
Mức độ tác động: **trung bình-cao**.
### 4. Warmup/scheduler chưa thực sự hoạt động
Trong [code/train.py](../code/train.py#L440)-[code/train.py](../code/train.py#L448):
- Nếu `--warmup`, LR của SAM được gán bằng `base_lr / warmup_period`, nhưng không có bước nào tăng dần lên `base_lr`.
- U-Net luôn dùng `base_lr`.
- `LambdaLR` được import tại [code/train.py](../code/train.py#L33) nhưng không dùng.
- Nhánh không `--AdamW` dùng `model.parameters()` tại [code/train.py](../code/train.py#L447), trong khi `model` không tồn tại trong scope này; chạy không có `--AdamW` sẽ lỗi.
Khuyến nghị:
- Sửa bug `model.parameters()` thành `SAM_model.parameters()`.
- Thêm scheduler theo iteration: warmup từ LR nhỏ lên target LR, sau đó cosine/poly decay.
- Tách LR cho SAM/LoRA và U-Net. LoRA/decoder của SAM thường nên có LR nhỏ hơn U-Net, hoặc ít nhất được log riêng.
- Log `sam_lr` và `unet_lr` mỗi epoch/iteration.
Mức độ tác động: **cao với cấu hình không AdamW**, **trung bình-cao với chất lượng train**.
### 5. Evaluation có nguy cơ load sai checkpoint
`train.py` tạo output tại [code/train.py](../code/train.py#L769):
```python
../model/<dataset>/train/<save_name>/
```
Nhưng `test.py` load tại [code/test.py](../code/test.py#L238) và [code/test.py](../code/test.py#L241):
```python
../model/<dataset>/<save_name>/...
```
Do đó, một run train bằng `--save_name prostate_dm1_lb40_full` sẽ cần test với `--save_name train/prostate_dm1_lb40_full`, không phải `--save_name prostate_dm1_lb40_full`. Đây là gotcha dễ đánh giá nhầm.
Thêm nữa, `test.py` default `--model unet` tại [code/test.py](../code/test.py#L26), trong khi `train.py` default `--model MedSAM` tại [code/train.py](../code/train.py#L47). Nếu người dùng truyền `--model MedSAM` vào `test.py`, code không vào nhánh `SAM` hoặc `unet` tại [code/test.py](../code/test.py#L237)-[code/test.py](../code/test.py#L242).
Khuyến nghị:
- Chuẩn hóa argument `--experiment_dir` hoặc `--checkpoint_dir`.
- Cho phép `--model SAM` và `--model MedSAM` cùng map vào nhánh LoRA SAM.
- Khi eval, assert checkpoint tồn tại và in đường dẫn checkpoint rõ ràng.
- Không tạo/chặn thư mục checkpoint trong eval như [code/test.py](../code/test.py#L295)-[code/test.py](../code/test.py#L298).
Mức độ tác động: **cao**, vì metric có thể không phản ánh run vừa train.
### 6. Metric HD/ASD xử lý mask rỗng chưa đúng
Trong [code/test.py](../code/test.py#L112)-[code/test.py](../code/test.py#L123) và bản copy trong [code/train.py](../code/train.py#L207)-[code/train.py](../code/train.py#L218), code chỉ check prediction rỗng:
- Pred rỗng: gán HD/ASD = 100.
- Pred không rỗng: gọi `medpy.binary.hd95/asd`.
Nhưng cần xử lý cả ground truth rỗng. Nếu pred và gt cùng rỗng, phạt 100 là không hợp lý. Nếu pred có object nhưng gt rỗng, MedPy có thể lỗi vì một input rỗng.
Khuyến nghị:
- `pred_empty and gt_empty`: Dice policy rõ ràng, HD/ASD = 0.
- `pred_empty xor gt_empty`: penalty cố định, ví dụ 100.
- Cả hai có object: gọi `hd95/asd`.
Mức độ tác động: **cao với dataset có slice/background rỗng**, đặc biệt BUSI/MNMS.
### 7. Dice metric có smoothing làm inflate kết quả object nhỏ
Trong [code/utils/metrics.py](../code/utils/metrics.py#L139)-[code/utils/metrics.py](../code/utils/metrics.py#L143), Dice dùng:
```python
(2 * intersection + 1.0) / (1.001 + segmentation_pixels + gt_label_pixels)
```
Với ground truth rỗng và prediction 1 pixel, Dice xấp xỉ `1 / 2.001`, tức gần 0.5. Điều này làm false positive nhỏ trong object nhỏ/rỗng nhìn tốt hơn thực tế.
Khuyến nghị:
- Dùng epsilon nhỏ, ví dụ `1e-6`, chỉ để tránh chia 0.
- Explicit policy:
- both empty -> 1.0 nếu coi là đúng background;
- one empty -> 0.0;
- both non-empty -> công thức Dice chuẩn.
- Báo cáo rõ policy trong paper/report để so sánh công bằng.
Mức độ tác động: **trung bình-cao**.
### 8. Trung bình metric theo batch thay vì theo ảnh
Trong evaluation, `dice_calcu` trả trung bình theo batch, sau đó cộng/chia theo `len(cur_dataloader)` tại [code/test.py](../code/test.py#L136)-[code/test.py](../code/test.py#L137). Nếu `test_bs > 1`, batch cuối nhỏ vẫn có trọng số bằng batch đầy.
Default `test_bs=1` nên hiện tại ít bị ảnh hưởng, nhưng code để người dùng đổi `--test_bs`.
Khuyến nghị:
- Tích lũy tổng metric theo số ảnh.
- Hoặc khóa `test_bs=1` khi tính metric paper.
- Log `num_cases`/`num_slices` mỗi domain.
Mức độ tác động: **trung bình**.
### 9. Labeled split lấy first-N sau khi sort file
Trong [code/train.py](../code/train.py#L415)-[code/train.py](../code/train.py#L422), labeled index là:
```python
lb_idxs = list(range(lb_num))
unlabeled_idxs = list(range(lb_num, data_num))
```
Trong dataloader, file được `imagelist.sort()` trước khi chọn index. Nếu tên file sắp theo patient/case/slice, labeled set sẽ không đại diện cho phân phối thật.
Khuyến nghị:
- Tạo split bằng RNG seed: shuffle index rồi lấy `lb_num`.
- Lưu split vào file JSON/TXT để reproduce.
- Với prostate/MNMS, nên split theo patient/case thay vì slice nếu có metadata.
- Với BUSI, nên stratified theo benign/malignant và kích thước lesion.
Mức độ tác động: **cao với hiệu năng mô hình và độ tin cậy thực nghiệm**.
### 10. Data transform có thể làm xấu ảnh đầu vào U-Net
Trong [code/dataloaders/custom_transforms.py](../code/dataloaders/custom_transforms.py#L766)-[code/dataloaders/custom_transforms.py](../code/dataloaders/custom_transforms.py#L778), `zoom(..., order=0)` được dùng cho cả `unet_size_img` và `unet_size_strong_aug`. `order=0` là nearest-neighbor, phù hợp cho label nhưng không phù hợp cho image liên tục.
Ảnh y tế bị resize nearest-neighbor có thể bị blocky/aliasing, làm U-Net học biên và texture kém hơn.
Khuyến nghị:
- Image: dùng bilinear (`order=1`) hoặc area/antialias.
- Label: giữ nearest-neighbor (`order=0`).
- Viết test shape/value range để đảm bảo không làm đổi label class.
Mức độ tác động: **trung bình**.
### 11. So sánh string bằng `is`
Trong [code/dataloaders/dataloader.py](../code/dataloaders/dataloader.py#L100), [code/dataloaders/dataloader.py](../code/dataloaders/dataloader.py#L231), [code/dataloaders/dataloader.py](../code/dataloaders/dataloader.py#L343) và nhiều vị trí khác, code dùng:
```python
if _img.mode is 'RGB':
```
`is` là identity comparison, không phải value comparison. Trong Python, string interning có thể làm nó "có vẻ đúng" trong một số trường hợp, nhưng không đảm bảo.
Khuyến nghị:
- Đổi tất cả thành `==`.
- Thêm lint rule hoặc test nhỏ cho dataloader mode conversion.
Mức độ tác động: **trung bình**.
### 12. SAM/MedSAM preprocessing cần được kiểm chứng bằng ablation
Khi tạo SAM tại [code/train.py](../code/train.py#L354)-[code/train.py](../code/train.py#L358), repo truyền:
```python
pixel_mean=[0, 0, 0], pixel_std=[1, 1, 1]
```
Trong khi `Normalize_tf(dataRange=[0,1])` đưa image về `[0,1]`. Nếu checkpoint là SAM gốc, pretrained normalization thường dựa trên pixel `[0,255]` và ImageNet mean/std trong [code/segment_anything/modeling/sam.py](../code/segment_anything/modeling/sam.py#L28)-[code/segment_anything/modeling/sam.py](../code/segment_anything/modeling/sam.py#L48). Nếu checkpoint MedSAM của dự án đã được train với `[0,1]` thì cách này có thể đúng.
Khuyến nghị:
- Chạy ablation:
- MedSAM checkpoint + current `[0,1]`.
- SAM checkpoint + SAM normalization gốc.
- MedSAM checkpoint + z-score/domain-wise normalization.
- Ghi rõ preprocessing theo checkpoint trong config/log.
Mức độ tác động: **cần kiểm chứng**, nhưng có thể lớn nếu normalization mismatch.
### 13. SAM đang train theo no-prompt semantic segmentation
Trong [code/segment_anything/modeling/sam.py](../code/segment_anything/modeling/sam.py#L64)-[code/segment_anything/modeling/sam.py](../code/segment_anything/modeling/sam.py#L66), `forward_train()` gọi prompt encoder với `points=None, boxes=None, masks=None`. Nghĩa là SAM/MedSAM được dùng như một semantic segmentation model không prompt.
Đây có thể là chủ đích của paper/repo, nhưng nó làm mất một lợi thế lớn của SAM: promptable segmentation.
Hướng phát triển:
- Sinh box prompt từ labeled mask cho supervised branch.
- Sinh box/point prompt từ pseudo-label có confidence cao cho unlabeled branch.
- Thêm random perturbation vào box prompt để tăng robust.
- Nếu tiếp tục no-prompt, nên đổi tên mode rõ ràng thành semantic mode và dọn API `multimask_output`.
Mức độ tác động: **trung bình-cao theo hướng research**, không nhất thiết là bug.
### 14. Resize relative position khi load checkpoint hard-code theo ViT-B
Trong [code/segment_anything/build_sam.py](../code/segment_anything/build_sam.py#L153), `global_rel_pos_keys` được chọn bằng substring `'2'`, `'5'`, `'8'`, `'11'`. Cách này đúng với ViT-B default indexes `[2, 5, 8, 11]`, nhưng sai với ViT-L/H và có nguy cơ match nhầm key.
Khuyến nghị:
- Truyền `encoder_global_attn_indexes` vào `load_from()`.
- Parse key dạng `image_encoder.blocks.{idx}.attn.rel_pos_h/w`.
- Resize bất cứ key relative-position nào có shape mismatch, thay vì hard-code theo substring.
Mức độ tác động: **cao nếu đổi `--vit_name` hoặc `--img_size`**.
### 15. Một số overhead trong SAM có thể giảm tốc độ train
Các điểm runtime:
- `PromptEncoder.get_dense_pe()` tạo positional encoding mới mỗi forward tại [code/segment_anything/modeling/prompt_encoder.py](../code/segment_anything/modeling/prompt_encoder.py#L62), được gọi trong [code/segment_anything/modeling/sam.py](../code/segment_anything/modeling/sam.py#L69).
- `get_rel_pos()` tạo `torch.arange` không chỉ định device tại [code/segment_anything/modeling/image_encoder.py](../code/segment_anything/modeling/image_encoder.py#L319)-[code/segment_anything/modeling/image_encoder.py](../code/segment_anything/modeling/image_encoder.py#L320).
- `unet_stu_output_ulb_x_w` tại [code/train.py](../code/train.py#L537) chỉ dùng để tính confidence, nhưng không bọc `torch.no_grad()`.
- Confidence/Dice trong training loop chuyển qua CPU/NumPy tại [code/train.py](../code/train.py#L542)-[code/train.py](../code/train.py#L546).
Khuyến nghị:
- Cache dense positional encoding theo device/dtype.
- Tạo `torch.arange(..., device=rel_pos.device)`.
- Bọc các forward/metric chỉ dùng cho confidence bằng `torch.no_grad()`.
- Đưa Dice confidence sang torch GPU, tránh `np.asarray(...cpu())` trong mỗi iteration.
Mức độ tác động: **trung bình về tốc độ**, có thể giúp train nhanh và ổn định hơn với batch lớn.
### 16. Checkpointing chưa đủ để resume/reproduce
Trong [code/train.py](../code/train.py#L729) và [code/train.py](../code/train.py#L753), repo chỉ lưu best student model weights. Không lưu optimizer, scheduler, scaler AMP, EMA teacher, iteration, RNG state. Ngoài ra `--save_model` mặc định tắt tại [code/train.py](../code/train.py#L78), nên người dùng có thể train xong mà không có checkpoint.
Khuyến nghị:
- Lưu checkpoint dict đầy đủ:
- `model_state_dict`
- `ema_model_state_dict`
- `optimizer_state_dict`
- `scaler_state_dict`
- `iter_num`, `epoch`, `best_metrics`
- RNG states
- Thêm `--resume`.
- Mặc định nên save checkpoint best, hoặc cảnh báo rõ nếu `--save_model` tắt.
Mức độ tác động: **trung bình về vận hành**, **cao nếu cần reproduce nghiên cứu**.
### 17. Những điểm maintainability nên dọn
- `icecream` import nhiều nơi nhưng không dùng: [code/segment_anything/modeling/sam.py](../code/segment_anything/modeling/sam.py#L10), [code/segment_anything/build_sam.py](../code/segment_anything/build_sam.py#L9), [code/sam_lora_image_encoder.py](../code/sam_lora_image_encoder.py#L14). Nên bỏ để giảm dependency.
- LoRA load checkpoint dùng `torch.load(filename)` không `map_location` tại [code/sam_lora_image_encoder.py](../code/sam_lora_image_encoder.py#L152), rồi gán `Parameter(saved_tensor)` tại [code/sam_lora_image_encoder.py](../code/sam_lora_image_encoder.py#L157). Nên load theo device của model và dùng `load_state_dict`/copy data có check missing keys.
- `MaskDecoder.forward()` bỏ qua `multimask_output` tại [code/segment_anything/modeling/mask_decoder.py](../code/segment_anything/modeling/mask_decoder.py#L102)-[code/segment_anything/modeling/mask_decoder.py](../code/segment_anything/modeling/mask_decoder.py#L111). Nếu dùng semantic class logits thì nên document/rename, nếu dùng SAM predictor thì nên khôi phục behavior.
- DSBN trong [code/networks/dsbn.py](../code/networks/dsbn.py#L26) chỉ lấy `domain_label[0]` cho cả batch. Nếu batch có nhiều domain, BN sai. Nên assert batch cùng domain hoặc split theo domain.
- WRN pooling trong [code/networks/wrn.py](../code/networks/wrn.py#L95) giả định feature map vuông. Nên dùng `F.adaptive_avg_pool2d(out, 1)`.
## Lộ trình tối ưu để phát triển dự án
### Giai đoạn 1: Sửa bug/metric để kết quả đáng tin
1. Sửa `DiceLossWithMask` confidence mask.
2. Sửa masked CE normalization.
3. Sửa `test.py` checkpoint path/model alias.
4. Sửa HD/ASD empty-mask handling.
5. Sửa Dice smoothing policy.
6. Sửa `model.parameters()` -> `SAM_model.parameters()` trong nhánh SGD.
Kỳ vọng: metric ổn định hơn, tránh đánh giá nhầm, unsupervised signal đúng hơn.
### Giai đoạn 2: Ổn định training
1. Đặt EMA teacher `eval()` khi tạo pseudo-label.
2. Thêm LR scheduler thật cho SAM và U-Net.
3. Log LR, mask ratio, supervised/unsupervised loss normalized.
4. Lưu split labeled bằng seed.
5. Lưu checkpoint dict đầy đủ và thêm resume.
Kỳ vọng: dễ reproduce, dễ debug, giảm dao động giữa seed/domain.
### Giai đoạn 3: Tối ưu tốc độ
1. Chuyển confidence Dice sang torch GPU.
2. Bọc các forward chỉ dùng để tính confidence bằng `no_grad()`.
3. Cache dense PE và relative-position indices trong SAM.
4. Tăng `num_workers`/`persistent_workers` nếu IO là bottleneck.
5. Bỏ import/dependency không dùng.
Kỳ vọng: train nhanh hơn, GPU ít idle hơn.
### Giai đoạn 4: Phát triển research/model
1. Thử prompt từ box/point lấy từ label và pseudo-label.
2. Thử uncertainty calibration thay vì threshold cố định 0.95.
3. Domain-aware sampling cho unlabeled data, tránh domain lớn lấn át.
4. Thêm test-time augmentation hoặc connected-component post-processing theo dataset.
5. Đánh giá EMA teacher như candidate checkpoint riêng.
Kỳ vọng: cải thiện khả năng tổng quát trong MiDSS/domain shift, đặc biệt khi labeled domain cách xa unlabeled domains.
## Các thí nghiệm nên chạy để xác nhận
Nên chạy ablation nhỏ trước khi đổi nhiều thứ cùng lúc:
| Thí nghiệm | Thay đổi | Metric cần xem |
| --- | --- | --- |
| Baseline | Code hiện tại, cùng seed | Dice từng domain, avg Dice, mask ratio |
| Loss-mask fix | Sửa Dice mask + masked CE | Dice, unsup loss scale, mask ratio |
| EMA eval | Teacher eval khi pseudo-label | Dice, self/mutual confidence |
| Scheduler | Warmup + cosine/poly decay | Loss curve, best Dice iter |
| Random split | Shuffle/stratified labeled split | Mean/std qua nhiều seed |
| SAM prompt | Box/point prompt supervised/unlabeled | SAM Dice, ensemble pseudo-label Dice |
Khi báo cáo, nên ghi rõ:
- Dataset, `lb_domain`, `lb_num`, seed.
- Checkpoint SAM/MedSAM và preprocessing.
- `save_name`, commit hash, command train/test.
- Best avg Dice và Dice từng class/domain.
- Có dùng EMA checkpoint hay student checkpoint.
## Kết luận
Dự án này có ý tưởng tốt: kết hợp conventional U-Net với foundation model SAM/MedSAM để hỗ trợ semi-supervised segmentation dưới domain shift. Tuy nhiên, hiện có một số điểm cài đặt có khả năng làm giảm kết quả hoặc làm metric không đáng tin, đặc biệt là **mask trong unsupervised Dice**, **normalization của masked CE**, **evaluation checkpoint path**, **empty-mask metric**, **EMA teacher mode**, và **split labeled data**.
Nếu chỉ có thời gian sửa ít, nên ưu tiên Giai đoạn 1 và Giai đoạn 2. Đây là nhóm thay đổi ít mang tính kiến trúc nhất, nhưng có khả năng tăng độ tin cậy và hiệu năng thực nghiệm rõ nhất.
|