update

Author: mjun0812

src/__init__.py

@@ -28,7 +28,6 @@
     build_dataloader,
     build_dataset,
     build_sampler,
-    is_ram_cache_supported,
     validate_ram_cache_config,
 )
 

src/dataloaders/__init__.py

@@ -7,7 +7,6 @@
     build_dataloader,
     build_dataset,
     build_sampler,
-    is_ram_cache_supported,
     validate_ram_cache_config,
 )
 from .tensor_cache import TensorCache

src/dataloaders/build.py

@@ -20,15 +20,6 @@
 """Registry for dataset classes."""
 
 
-def is_ram_cache_supported() -> bool:
-    """Check if RAM cache is supported on the current platform.
-
-    Returns:
-        bool: True if RAM cache is supported (Linux only), False otherwise.
-    """
-    return platform.system() == "Linux"
-
-
 def validate_ram_cache_config(use_ram_cache: bool, verbose: bool = True) -> bool:
     """Validate RAM cache configuration for current platform.
 
@@ -39,7 +30,7 @@ def validate_ram_cache_config(use_ram_cache: bool, verbose: bool = True) -> bool
     Returns:
         bool: True if RAM cache can be used, False otherwise.
     """
-    if use_ram_cache and not is_ram_cache_supported():
+    if use_ram_cache and platform.system() != "Linux":
         if verbose:
             logger.warning(
                 f"RAM cache is only supported on Linux. Current OS: {platform.system()}. "
@@ -161,7 +152,7 @@ def build_sampler(
     return sampler, batch_sampler
 
 
-def collate(batch: list[DatasetOutput]) -> dict[str, torch.Tensor]:
+def collate(batch: list[DatasetOutput]) -> DatasetOutput:
     """Collate function for combining dataset outputs into batches.
 
     Automatically handles stacking of tensors and conversion of numeric
@@ -183,4 +174,4 @@ def collate(batch: list[DatasetOutput]) -> dict[str, torch.Tensor]:
             output[k] = torch.stack(output[k], dim=0)
         elif isinstance(output[k][0], (int, float)):
             output[k] = torch.tensor(output[k])
-    return output
+    return DatasetOutput(**output)

src/dataloaders/dummy.py

@@ -1,6 +1,5 @@
 from dataclasses import dataclass
 
-import torch
 from torchvision.transforms.v2 import Compose
 
 from .base import BaseDataset
@@ -58,4 +57,4 @@ def __getitem__(self, idx: int) -> DatasetOutput:
         Returns:
             Synthetic data sample with ones tensor and modulo label.
         """
-        return DatasetOutput(data=torch.ones(8), label=idx % 4)
+        return DatasetOutput.dummy()

src/dataloaders/types.py

@@ -1,5 +1,6 @@
 from typing import NotRequired, TypedDict
 
+import torch
 from torch import Tensor
 
 
@@ -11,5 +12,24 @@ class DatasetOutput(TypedDict, total=False):
         label: The target label as integer (optional).
     """
 
+    # GPUに転送するキーのリスト
+    __gpu_keys__: tuple[str] = ("data", "label")
+
     data: NotRequired[Tensor]
     label: NotRequired[int]
+
+    @classmethod
+    def dummy(cls, batch: int | None = None) -> "DatasetOutput":
+        """Create dummy DatasetOutput for testing.
+
+        Args:
+            batch: Batch size. If None, returns single sample with label as int.
+                   If specified, returns batched data with label as Tensor.
+
+        Returns:
+            DatasetOutput with dummy values.
+        """
+        if batch is None:
+            return cls(data=torch.randn(3, 224, 224), label=torch.randint(0, 1000, (1,)).item())
+        else:
+            return cls(data=torch.randn(batch, 3, 224, 224), label=torch.randint(0, 1000, (batch,)))

src/models/types.py

@@ -1,5 +1,6 @@
 from typing import NotRequired, Required, TypedDict
 
+import torch
 from torch import Tensor
 
 
@@ -12,6 +13,26 @@ class LossOutput(TypedDict, total=False):
 
     total_loss: Required[Tensor]
 
+    @classmethod
+    def dummy(cls) -> "LossOutput":
+        """Create dummy LossOutput for testing."""
+        return cls(total_loss=torch.rand(1))
+
+
+class PredOutput(TypedDict, total=False):
+    """Output structure for model predictions.
+
+    Attributes:
+        preds: Dictionary containing model predictions.
+    """
+
+    preds: NotRequired[Tensor]
+
+    @classmethod
+    def dummy(cls, batch: int = 1) -> "PredOutput":
+        """Create dummy PredOutput for testing."""
+        return cls(preds=torch.rand(batch, 10))
+
 
 class ModelOutput(TypedDict, total=False):
     """Output structure for model forward passes.
@@ -22,4 +43,9 @@ class ModelOutput(TypedDict, total=False):
     """
 
     losses: NotRequired[LossOutput]
-    preds: NotRequired[dict[str, Tensor]]
+    preds: NotRequired[PredOutput]
+
+    @classmethod
+    def dummy(cls, batch: int = 1) -> "ModelOutput":
+        """Create dummy ModelOutput for testing."""
+        return cls(losses=LossOutput.dummy(), preds=PredOutput.dummy(batch))

src/runner/base_trainer.py

@@ -278,7 +278,7 @@ def _do_one_epoch(self, state: TrainingState) -> EpochResult:
             with torch.set_grad_enabled(state.phase == "train"):
                 # Prepare Input
                 for k, v in inputs.items():
-                    if isinstance(v, torch.Tensor):
+                    if isinstance(v, torch.Tensor) and k in DatasetOutput.__gpu_keys__:
                         inputs[k] = v.to(self.device, non_blocking=True)
                 if self.batched_transforms[state.phase]:
                     inputs = self.batched_transforms[state.phase](inputs)

tests/test_evaluator.py

@@ -1,13 +1,13 @@
 from pathlib import Path
 
 import pytest
-import torch
 import yaml
 from torchmetrics import MetricCollection
 
 from src.config.config import EvaluatorConfig
 from src.dataloaders.types import DatasetOutput
 from src.evaluator.build import EVALUATOR_REGISTRY
+from src.models.types import PredOutput
 
 CONFIG_PATH = "config/__base__/evaluator/dummy_evaluator.yaml"
 
@@ -17,18 +17,13 @@ class TestEvaluatorHelpers:
 
     @staticmethod
     def create_dataset_data(batch_size: int = 2) -> DatasetOutput:
-        """Create mock HODatasetData for testing"""
-        data = DatasetOutput()
-        data["data"] = torch.rand(batch_size, 3, 224, 224)
-        data["label"] = torch.randint(0, 10, (batch_size,))
-        return data
+        """Create mock dataset data for testing"""
+        return DatasetOutput.dummy(batch=batch_size)
 
     @staticmethod
-    def create_pred_output(batch_size: int = 2) -> dict:
-        """Create mock HOPEPredOutput for testing"""
-        output = {}
-        output["pred"] = torch.rand(batch_size, 10)
-        return output
+    def create_pred_output(batch_size: int = 2) -> PredOutput:
+        """Create mock pred output for testing"""
+        return PredOutput.dummy(batch=batch_size)
 
 
 class TestEvaluatorFromConfig:

tests/test_modules.py

@@ -123,7 +123,9 @@ def test_dataloader() -> None:
         transform = build_transforms(dataset_cfg.test.transforms)
         dataset = build_dataset(dataset_cfg.test, transform)
         _, batch_sampler = build_sampler(dataset, phase="test", batch_size=2)
-        dataloader = build_dataloader(dataset, num_workers=2, batch_sampler=batch_sampler)
+        dataloader = build_dataloader(
+            dataset, num_workers=2, batch_sampler=batch_sampler, pin_memory=False
+        )
         if dataset_cfg.test.batch_transforms is not None:
             batched_transform = build_batched_transform(dataset_cfg.test.batch_transforms)
         else: