update

Author: mjun0812

src/dataloaders/types.py

@@ -1,3 +1,4 @@
+import random
 from typing import NotRequired, TypedDict
 
 import torch
@@ -30,6 +31,8 @@ def dummy(cls, batch: int | None = None) -> "DatasetOutput":
             DatasetOutput with dummy values.
         """
         if batch is None:
-            return cls(data=torch.randn(3, 224, 224), label=torch.randint(0, 1000, (1,)).item())
+            return cls(data=torch.randn(8), label=random.randint(0, 1000))
         else:
-            return cls(data=torch.randn(batch, 3, 224, 224), label=torch.randint(0, 1000, (batch,)))
+            return cls(
+                data=torch.randn(batch, 8), label=[random.randint(0, 1000) for _ in range(batch)]
+            )

src/runner/base_trainer.py

@@ -427,14 +427,38 @@ def save_state(self, epoch: int, output_dir: Path) -> dict[str, Path]:
             create_symlink(weight_path, final_model_path)
 
             save_optimizer(self.optimizer, optimizer_path)
+            # 保存したoptimizer_path以外のoptimizerファイルを削除
+            optimizer_dir = Path(output_dir) / "optimizers"
+            for optimizer_file in optimizer_dir.glob("optimizer_epoch_*.pth"):
+                if (
+                    str(optimizer_file) != optimizer_path
+                    and str(optimizer_file) != final_optimizer_path
+                ):
+                    optimizer_file.unlink()
             create_symlink(optimizer_path, final_optimizer_path)
 
             if self.epoch_lr_scheduler:
                 save_lr_scheduler(self.epoch_lr_scheduler, epoch_scheduler_path)
+                # 保存したschedulerファイル以外のschedulerファイルを削除
+                scheduler_dir = Path(output_dir) / "schedulers"
+                for scheduler_file in scheduler_dir.glob("epoch_scheduler_epoch_*.pth"):
+                    if (
+                        str(scheduler_file) != epoch_scheduler_path
+                        and str(scheduler_file) != final_epoch_scheduler_path
+                    ):
+                        scheduler_file.unlink()
                 create_symlink(epoch_scheduler_path, final_epoch_scheduler_path)
 
             if self.iter_lr_scheduler:
                 save_lr_scheduler(self.iter_lr_scheduler, iter_scheduler_path)
+                # 保存したschedulerファイル以外のschedulerファイルを削除
+                scheduler_dir = Path(output_dir) / "schedulers"
+                for scheduler_file in scheduler_dir.glob("iter_scheduler_epoch_*.pth"):
+                    if (
+                        str(scheduler_file) != iter_scheduler_path
+                        and str(scheduler_file) != final_iter_scheduler_path
+                    ):
+                        scheduler_file.unlink()
                 create_symlink(iter_scheduler_path, final_iter_scheduler_path)
 
     def load_state(

test.py

@@ -29,7 +29,7 @@ def do_test(
     logger.info("Loading Dataset...")
     transform = build_transforms(cfg.dataset.test.transforms)
     if cfg.dataset.test.batch_transforms is not None:
-        batched_transform = build_batched_transform(cfg.dataset.test.batch_transforms)
+        batched_transform = build_batched_transform(cfg.dataset.test.batch_transforms).to(device)
     else:
         batched_transform = None
     dataset = build_dataset(cfg.dataset.test, transform)