konfai 1.1.1__py3-none-any.whl → 1.1.3__py3-none-any.whl

konfai/data/transform.py

@@ -6,7 +6,7 @@ from abc import ABC, abstractmethod
 import torch.nn.functional as F
 from typing import Any, Union
 
-from konfai.utils.utils import _getModule, NeedDevice, _resample_affine, _affine_matrix
+from konfai.utils.utils import _getModule, NeedDevice, _resample_affine, _affine_matrix, TransformError
 from konfai.utils.dataset import Dataset, Attribute, data_to_image, image_to_data
 from konfai.utils.config import config
 
@@ -52,7 +52,7 @@ class Clip(Transform):
         input[torch.where(input < self.min_value)] = self.min_value
         input[torch.where(input > self.max_value)] = self.max_value
         if self.saveClip_min:
-            cache_attribute["Min"] = self .min_value
+            cache_attribute["Min"] = self.min_value
         if self.saveClip_max:
             cache_attribute["Max"] = self.max_value
         return input
@@ -211,39 +211,67 @@ class Resample(Transform, ABC):
         _ = cache_attribute.pop_np_array("Spacing")
         return self._resample(input, [int(size) for size in size_1])
 
-class ResampleIsotropic(Resample):
+class ResampleToResolution(Resample):
+
+    def __init__(self, spacing : list[Union[float, None]] = [1., 1., 1.]) -> None:
+        self.spacing = torch.tensor([0 if s < 0 else s for s in spacing])
 
-    def __init__(self, spacing : list[float] = [1., 1., 1.]) -> None:
-        self.spacing = torch.tensor(spacing, dtype=torch.float64)
-        
     def transformShape(self, shape: list[int], cache_attribute: Attribute) -> list[int]:
-        assert "Spacing" in cache_attribute, "Error no spacing"
-        resize_factor = self.spacing/cache_attribute.get_tensor("Spacing").flip(0)
-        return  [int(x) for x in (torch.tensor(shape) * 1/resize_factor)]
+        if "Spacing" not in cache_attribute:
+            TransformError("Missing 'Spacing' in cache attributes, the data is likely not a valid image.",
+                        "Make sure your input is a image (e.g., .nii, .mha) with proper metadata.")
+        if len(shape) != len(self.spacing):
+            TransformError("Shape and spacing dimensions do not match: shape={shape}, spacing={self.spacing}")
+        image_spacing = cache_attribute.get_tensor("Spacing").flip(0)
+        spacing = self.spacing
+        
+        for i, s in enumerate(self.spacing):
+            if s == 0:
+                spacing[i] = image_spacing[i]
+        resize_factor = spacing/cache_attribute.get_tensor("Spacing").flip(0)
+        return [int(x) for x in (torch.tensor(shape) * 1/resize_factor)]
 
     def __call__(self, name: str, input : torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:
-        assert "Spacing" in cache_attribute, "Error no spacing"
-        resize_factor = self.spacing/cache_attribute.get_tensor("Spacing").flip(0)
-        cache_attribute["Spacing"] = self.spacing.flip(0)
+        image_spacing = cache_attribute.get_tensor("Spacing").flip(0)
+        spacing = self.spacing
+        for i, s in enumerate(self.spacing):
+            if s == 0:
+                spacing[i] = image_spacing[i]
+        resize_factor = spacing/cache_attribute.get_tensor("Spacing").flip(0)
+        cache_attribute["Spacing"] = spacing.flip(0)
         cache_attribute["Size"] = np.asarray([int(x) for x in torch.tensor(input.shape[1:])])
         size = [int(x) for x in (torch.tensor(input.shape[1:]) * 1/resize_factor)]
         cache_attribute["Size"] = np.asarray(size)
         return self._resample(input, size)
 
-class ResampleResize(Resample):
+class ResampleToSize(Resample):
 
     def __init__(self, size : list[int] = [100,512,512]) -> None:
         self.size = size
 
     def transformShape(self, shape: list[int], cache_attribute: Attribute) -> list[int]:
-        return self.size
+        if "Spacing" not in cache_attribute:
+            TransformError("Missing 'Spacing' in cache attributes, the data is likely not a valid image.",
+                        "Make sure your input is a image (e.g., .nii, .mha) with proper metadata.")
+        if len(shape) != len(self.size):
+            TransformError("Shape and spacing dimensions do not match: shape={shape}, spacing={self.spacing}")
+        size = self.size
+        for i, s in enumerate(self.size):
+            if s == -1:
+                size[i] = shape[i]
+        return size
     
     def __call__(self, name: str, input: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:
+        size = self.size
+        image_size =  np.asarray([int(x) for x in torch.tensor(input.shape[1:])])
+        for i, s in enumerate(self.size):
+            if s is None:
+                size[i] = image_size[i]
         if "Spacing" in cache_attribute:
-            cache_attribute["Spacing"] = torch.flip(torch.tensor(list(input.shape[1:]))/torch.tensor(self.size)*torch.flip(cache_attribute.get_tensor("Spacing"), dims=[0]), dims=[0])
-        cache_attribute["Size"] = np.asarray([int(x) for x in torch.tensor(input.shape[1:])])
-        cache_attribute["Size"] = self.size
-        return self._resample(input, self.size)
+            cache_attribute["Spacing"] = torch.flip(torch.tensor(image_size)/torch.tensor(size)*torch.flip(cache_attribute.get_tensor("Spacing"), dims=[0]), dims=[0])
+        cache_attribute["Size"] = image_size
+        cache_attribute["Size"] = size
+        return self._resample(input, size)
 
 class ResampleTransform(Transform):
 
@@ -412,8 +440,8 @@ class FlatLabel(Transform):
 
 class Save(Transform):
 
-    def __init__(self, save: str) -> None:
-        self.save = save
+    def __init__(self, dataset: str) -> None:
+        self.dataset = dataset
     
     def __call__(self, name: str, input : torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:
         return input
@@ -528,7 +556,7 @@ class OneHot(Transform):
         self.num_classes = num_classes
 
     def __call__(self, name: str, input: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:
-        result =  F.one_hot(input.type(torch.int64), num_classes=self.num_classes).permute(0, len(input.shape), *[i+1 for i in range(len(input.shape)-1)]).float().squeeze(2)
+        result =  F.one_hot(input.type(torch.int64), num_classes=self.num_classes).permute(0, len(input.shape), *[i+1 for i in range(len(input.shape)-1)]).float().squeeze(0)
         return result
     
     def inverse(self, name: str, input: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:

konfai/evaluator.py

@@ -9,7 +9,7 @@ import builtins
 import importlib
 from konfai import EVALUATIONS_DIRECTORY, PREDICTIONS_DIRECTORY, KONFAI_ROOT, CONFIG_FILE
 from konfai.utils.config import config
-from konfai.utils.utils import _getModule, DistributedObject, synchronize_data
+from konfai.utils.utils import _getModule, DistributedObject, synchronize_data, EvaluatorError
 from konfai.data.data_manager import DataMetric
 
 class CriterionsAttr():
@@ -54,7 +54,8 @@ class Statistics():
             if name_dataset not in self.measures:
                 self.measures[name_dataset] = {}
             self.measures[name_dataset][name] = value
-            
+    
+    @staticmethod   
     def getStatistic(values: list[float]) -> dict[str, float]:
         return {"max": np.max(values), "min": np.min(values), "std": np.std(values), "25pc": np.percentile(values, 25), "50pc": np.percentile(values, 50), "75pc": np.percentile(values, 75), "mean": np.mean(values), "count": len(values)}
     
@@ -91,7 +92,6 @@ class Evaluator(DistributedObject):
             exit(0)
         super().__init__(train_name)
         self.metric_path = EVALUATIONS_DIRECTORY()+self.name+"/"
-        self.predict_path = PREDICTIONS_DIRECTORY()+self.name+"/"
         self.metricsLoader = metrics
         self.dataset = dataset
         self.metrics = {k: v.getTargetsCriterions(k) for k, v in self.metricsLoader.items()}
@@ -102,10 +102,10 @@ class Evaluator(DistributedObject):
         result = {}
         for output_group in self.metrics:
             for target_group in self.metrics[output_group]:
-                targets = [data_dict[group][0] for group in target_group.split("/") if group in data_dict]
+                targets = [data_dict[group][0].to(0) if torch.cuda.is_available() else data_dict[group][0] for group in target_group.split("/") if group in data_dict]
                 name = data_dict[output_group][1][0]
                 for metric in self.metrics[output_group][target_group]:
-                    result["{}:{}:{}".format(output_group, target_group, metric.__class__.__name__)] = metric(data_dict[output_group][0], *targets).item()
+                    result["{}:{}:{}".format(output_group, target_group, metric.__class__.__name__)] = metric(data_dict[output_group][0].to(0) if torch.cuda.is_available() else data_dict[output_group][0], *targets).item()
         statistics.add(result, name)
         return result
     
@@ -126,9 +126,26 @@ class Evaluator(DistributedObject):
 
         self.dataloader = self.dataset.getData(world_size)
 
+        groupsDest = [group for groups in self.dataset.groups_src.values() for group in groups]
+
+        missing_outputs = set(self.metrics.keys()) - set(groupsDest)
+        if missing_outputs:
+            raise EvaluatorError(
+                f"The following metric output groups are missing from 'groupsDest': {sorted(missing_outputs)}. ",
+                f"Available groups: {sorted(groupsDest)}"
+            )
+
+        target_groups = {target for targets in self.metrics.values() for target in targets}
+        missing_targets = target_groups - set(groupsDest)
+        if missing_targets:
+            raise EvaluatorError(
+                f"The following metric target groups are missing from 'groupsDest': {sorted(missing_targets)}. ",
+                f"Available groups: {sorted(groupsDest)}"
+            )
+
     def run_process(self, world_size: int, global_rank: int, gpu: int, dataloaders: list[DataLoader]):
         description = lambda measure : "Metric TRAIN : {} ".format(" | ".join("{}: {:.2f}".format(k, v) for k, v in measure.items()) if measure is not None else "")
-        with tqdm.tqdm(iterable = enumerate(dataloaders[0]), leave=False, desc = description(None), total=len(dataloaders[0])) as batch_iter:
+        with tqdm.tqdm(iterable = enumerate(dataloaders[0]), leave=True, desc = description(None), total=len(dataloaders[0]), ncols=0) as batch_iter:
             for _, data_dict in batch_iter:
                 batch_iter.set_description(description(self.update({k: (v[0], v[4]) for k,v in data_dict.items()}, self.statistics_train)))
         outputs = synchronize_data(world_size, gpu, self.statistics_train.measures)
@@ -136,7 +153,7 @@ class Evaluator(DistributedObject):
             self.statistics_train.write(outputs)
         if len(dataloaders) == 2:
             description = lambda measure : "Metric VALIDATION : {} ".format(" | ".join("{}: {:.2f}".format(k, v) for k, v in measure.items()) if measure is not None else "")
-            with tqdm.tqdm(iterable = enumerate(dataloaders[1]), leave=False, desc = description(None), total=len(dataloaders[1])) as batch_iter:
+            with tqdm.tqdm(iterable = enumerate(dataloaders[1]), leave=True, desc = description(None), total=len(dataloaders[1]), ncols=0) as batch_iter:
                 for _, data_dict in batch_iter:
                     batch_iter.set_description(description(self.update({k: (v[0], v[4]) for k,v in data_dict.items()}, self.statistics_validation)))
             outputs = synchronize_data(world_size, gpu, self.statistics_validation.measures)

konfai/main.py

@@ -3,25 +3,27 @@ import os
 from torch.cuda import device_count
 import torch.multiprocessing as mp
 from konfai.utils.utils import setup, TensorBoard, Log
+from konfai import KONFAI_NB_CORES
 
 import sys
 sys.path.insert(0, os.getcwd())
 
 def main():
+    import tracemalloc
+
     parser = argparse.ArgumentParser(description="KonfAI", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     try:
         with setup(parser) as distributedObject:
-            with Log(distributedObject.name):
+            with Log(distributedObject.name, 0):
                 world_size = device_count()
                 if world_size == 0:
-                    world_size = 1
+                    world_size = int(KONFAI_NB_CORES())
                 distributedObject.setup(world_size)
                 with TensorBoard(distributedObject.name):
                     mp.spawn(distributedObject, nprocs=world_size)
     except KeyboardInterrupt:
         print("\n[KonfAI] Manual interruption (Ctrl+C)")
 
-
 def cluster():
     parser = argparse.ArgumentParser(description="KonfAI", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
 

konfai/models/segmentation/UNet.py

@@ -7,32 +7,33 @@ from konfai.data.patching import ModelPatch
 
 class UNetHead(network.ModuleArgsDict):
 
-    def __init__(self, in_channels: int, nb_class: int, dim: int) -> None:
+    def __init__(self, in_channels: int, nb_class: int, dim: int, level: int) -> None:
         super().__init__()
         self.add_module("Conv", blocks.getTorchModule("Conv", dim)(in_channels = in_channels, out_channels = nb_class, kernel_size = 1, stride = 1, padding = 0))
         self.add_module("Softmax", torch.nn.Softmax(dim=1))
         self.add_module("Argmax", blocks.ArgMax(dim=1))
+        
 
 class UNetBlock(network.ModuleArgsDict):
 
-    def __init__(self, channels: list[int], nb_conv_per_stage: int, blockConfig: blocks.BlockConfig, downSampleMode: blocks.DownSampleMode, upSampleMode: blocks.UpSampleMode, attention : bool, block: type, nb_class: int, dim: int, i : int = 0, mri: bool = False) -> None:
+    def __init__(self, channels: list[int], nb_conv_per_stage: int, blockConfig: blocks.BlockConfig, downSampleMode: blocks.DownSampleMode, upSampleMode: blocks.UpSampleMode, attention : bool, block: type, nb_class: int, dim: int, i : int = 0) -> None:
         super().__init__()
         blockConfig_stride = blockConfig
         if i > 0:
             if downSampleMode != blocks.DownSampleMode.CONV_STRIDE:
                 self.add_module(downSampleMode.name, blocks.downSample(in_channels=channels[0], out_channels=channels[1], downSampleMode=downSampleMode, dim=dim))
             else:
-                blockConfig_stride = blocks.BlockConfig(blockConfig.kernel_size, (1,2,2) if mri and i > 4 else 2, blockConfig.padding, blockConfig.bias, blockConfig.activation, blockConfig.normMode)
+                blockConfig_stride = blocks.BlockConfig(blockConfig.kernel_size, 2, blockConfig.padding, blockConfig.bias, blockConfig.activation, blockConfig.normMode)
         self.add_module("DownConvBlock", block(in_channels=channels[0], out_channels=channels[1], blockConfigs=[blockConfig_stride]+[blockConfig]*(nb_conv_per_stage-1), dim=dim))
         if len(channels) > 2:
-            self.add_module("UNetBlock_{}".format(i+1), UNetBlock(channels[1:], nb_conv_per_stage, blockConfig, downSampleMode, upSampleMode, attention, block, nb_class, dim, i+1, mri=mri))
+            self.add_module("UNetBlock_{}".format(i+1), UNetBlock(channels[1:], nb_conv_per_stage, blockConfig, downSampleMode, upSampleMode, attention, block, nb_class, dim, i+1))
             self.add_module("UpConvBlock", block(in_channels=(channels[1]+channels[2]) if upSampleMode != blocks.UpSampleMode.CONV_TRANSPOSE else channels[1]*2, out_channels=channels[1], blockConfigs=[blockConfig]*nb_conv_per_stage, dim=dim))
             if nb_class > 0:
-                self.add_module("Head", UNetHead(channels[1], nb_class, dim), out_branch=[-1])
+                self.add_module("Head", UNetHead(channels[1], nb_class, dim, i), out_branch=[-1])
         if i > 0:
             if attention:
                 self.add_module("Attention", blocks.Attention(F_g=channels[1], F_l=channels[0], F_int=channels[0], dim=dim), in_branch=[1, 0], out_branch=[1])
-            self.add_module(upSampleMode.name, blocks.upSample(in_channels=channels[1], out_channels=channels[0], upSampleMode=upSampleMode, dim=dim, kernel_size=(1,2,2) if mri and i > 4 else 2, stride=(1,2,2) if mri and i > 4 else 2))
+            self.add_module(upSampleMode.name, blocks.upSample(in_channels=channels[1], out_channels=channels[0], upSampleMode=upSampleMode, dim=dim, kernel_size=2, stride=2))
             self.add_module("SkipConnection", blocks.Concat(), in_branch=[0, 1])
 
 class UNet(network.Network):
@@ -51,8 +52,6 @@ class UNet(network.Network):
                     downSampleMode: str = "MAXPOOL",
                     upSampleMode: str = "CONV_TRANSPOSE",
                     attention : bool = False,
-                    blockType: str = "Conv",
-                    mri: bool = False) -> None:
+                    blockType: str = "Conv") -> None:
         super().__init__(in_channels = channels[0], optimizer = optimizer, schedulers = schedulers, outputsCriterions = outputsCriterions, patch=patch, dim = dim)        
-        self.add_module("UNetBlock_0", UNetBlock(channels, nb_conv_per_stage, blockConfig, downSampleMode=blocks.DownSampleMode._member_map_[downSampleMode], upSampleMode=blocks.UpSampleMode._member_map_[upSampleMode], attention=attention, block = blocks.ConvBlock if blockType == "Conv" else blocks.ResBlock, nb_class=nb_class, dim=dim, mri = mri))    
-        
+        self.add_module("UNetBlock_0", UNetBlock(channels, nb_conv_per_stage, blockConfig, downSampleMode=blocks.DownSampleMode._member_map_[downSampleMode], upSampleMode=blocks.UpSampleMode._member_map_[upSampleMode], attention=attention, block = blocks.ConvBlock if blockType == "Conv" else blocks.ResBlock, nb_class=nb_class, dim=dim))    

konfai/network/network.py

@@ -39,7 +39,6 @@ class OptimizerLoader():
         self.name = name
     
     def getOptimizer(self, key: str, parameter: Iterator[torch.nn.parameter.Parameter]) -> torch.optim.Optimizer:
-        torch.optim.AdamW
         return config("{}.Model.{}.Optimizer".format(KONFAI_ROOT(), key))(getattr(importlib.import_module('torch.optim'), self.name))(parameter, config = None)
         
 class SchedulerStep():

konfai/predictor.py

@@ -91,13 +91,30 @@ class OutDataset(Dataset, NeedDevice, ABC):
         super().write(self.group, name, layer.numpy(), self.attributes[index][0][0])
         self.attributes.pop(index)
 
+class Reduction():
+
+    def __init__(self):
+        pass
+
+class ReductionMean():
+
+    def __init__(self):
+        pass
+    
+class ReductionMedian():
+
+    def __init__(self):
+        pass
+    
+    
+
 class OutSameAsGroupDataset(OutDataset):
 
     @config("OutDataset")
-    def __init__(self, dataset_filename: str = "./Dataset:mha", group: str = "default", sameAsGroup: str = "default", pre_transforms : dict[str, TransformLoader] = {"default:Normalize": TransformLoader()}, post_transforms : dict[str, TransformLoader] = {"default:Normalize": TransformLoader()}, final_transforms : dict[str, TransformLoader] = {"default:Normalize": TransformLoader()}, patchCombine: Union[str, None] = None, redution: str = "mean", inverse_transform: bool = True) -> None:
+    def __init__(self, dataset_filename: str = "./Dataset:mha", group: str = "default", sameAsGroup: str = "default", pre_transforms : dict[str, TransformLoader] = {"default:Normalize": TransformLoader()}, post_transforms : dict[str, TransformLoader] = {"default:Normalize": TransformLoader()}, final_transforms : dict[str, TransformLoader] = {"default:Normalize": TransformLoader()}, patchCombine: Union[str, None] = None, reduction: str = "mean", inverse_transform: bool = True) -> None:
         super().__init__(dataset_filename, group, pre_transforms, post_transforms, final_transforms, patchCombine)
         self.group_src, self.group_dest = sameAsGroup.split(":")
-        self.redution = redution
+        self.reduction = reduction
         self.inverse_transform = inverse_transform
 
     def addLayer(self, index_dataset: int, index_augmentation: int, index_patch: int, layer: torch.Tensor, dataset: DatasetIter):
@@ -151,9 +168,9 @@ class OutSameAsGroupDataset(OutDataset):
         self.output_layer_accumulator.pop(index)
         dtype = result.dtype
 
-        if self.redution == "mean":
+        if self.reduction == "mean":
             result = torch.mean(result.float(), dim=0).to(dtype)
-        elif self.redution == "median":
+        elif self.reduction == "median":
             result, _ = torch.median(result.float(), dim=0)
         else:
             raise NameError("Reduction method does not exist (mean, median)")
@@ -201,7 +218,7 @@ class OutDatasetLoader():
 
 class _Predictor():
 
-    def __init__(self, world_size: int, global_rank: int, local_rank: int, predict_path: str, data_log: Union[list[str], None], outsDataset: dict[str, OutDataset], modelComposite: DDP, dataloader_prediction: DataLoader) -> None:
+    def __init__(self, world_size: int, global_rank: int, local_rank: int, autocast: bool, predict_path: str, data_log: Union[list[str], None], outsDataset: dict[str, OutDataset], modelComposite: DDP, dataloader_prediction: DataLoader) -> None:
         self.world_size = world_size        
         self.global_rank = global_rank
         self.local_rank = local_rank
@@ -209,7 +226,7 @@ class _Predictor():
         self.modelComposite = modelComposite
         self.dataloader_prediction = dataloader_prediction
         self.outsDataset = outsDataset
-
+        self.autocast = autocast
         
         self.it = 0
 
@@ -239,21 +256,22 @@ class _Predictor():
         self.modelComposite.eval()  
         self.modelComposite.module.setState(NetState.PREDICTION)
         desc = lambda : "Prediction : {}".format(description(self.modelComposite))
-        self.dataloader_prediction.dataset.load()
-        with tqdm.tqdm(iterable = enumerate(self.dataloader_prediction), leave=False, desc = desc(), total=len(self.dataloader_prediction), disable=self.global_rank != 0 and "KONFAI_CLUSTER" not in os.environ) as batch_iter:
+        self.dataloader_prediction.dataset.load("Prediction")
+        with tqdm.tqdm(iterable = enumerate(self.dataloader_prediction), leave=True, desc = desc(), total=len(self.dataloader_prediction), ncols=0) as batch_iter:
             dist.barrier()
             for it, data_dict in batch_iter:
-                input = self.getInput(data_dict)
-                for name, output in self.modelComposite(input, list(self.outsDataset.keys())):
-                    self._predict_log(data_dict)
-                    outDataset = self.outsDataset[name]
-                    for i, (index, patch_augmentation, patch_index) in enumerate([(int(index), int(patch_augmentation), int(patch_index)) for index, patch_augmentation, patch_index in zip(list(data_dict.values())[0][1], list(data_dict.values())[0][2], list(data_dict.values())[0][3])]):
-                        outDataset.addLayer(index, patch_augmentation, patch_index, output[i].cpu(), self.dataset)
-                        if outDataset.isDone(index):
-                            outDataset.write(index, self.dataset.getDatasetFromIndex(list(data_dict.keys())[0], index).name.split("/")[-1], outDataset.getOutput(index, self.dataset))
-
-                batch_iter.set_description(desc())
-                self.it += 1
+                with torch.amp.autocast('cuda', enabled=self.autocast):
+                    input = self.getInput(data_dict)
+                    for name, output in self.modelComposite(input, list(self.outsDataset.keys())):
+                        self._predict_log(data_dict)
+                        outDataset = self.outsDataset[name]
+                        for i, (index, patch_augmentation, patch_index) in enumerate([(int(index), int(patch_augmentation), int(patch_index)) for index, patch_augmentation, patch_index in zip(list(data_dict.values())[0][1], list(data_dict.values())[0][2], list(data_dict.values())[0][3])]):
+                            outDataset.addLayer(index, patch_augmentation, patch_index, output[i].cpu(), self.dataset)
+                            if outDataset.isDone(index):
+                                outDataset.write(index, self.dataset.getDatasetFromIndex(list(data_dict.keys())[0], index).name.split("/")[-1], outDataset.getOutput(index, self.dataset))
+
+                    batch_iter.set_description(desc())
+                    self.it += 1
 
     def _predict_log(self, data_dict : dict[str, tuple[torch.Tensor, int, int, int]]):
         measures = DistributedObject.getMeasure(self.world_size, self.global_rank, self.local_rank, {"" : self.modelComposite.module}, 1)
@@ -320,6 +338,7 @@ class Predictor(DistributedObject):
                     train_name: str = "name",
                     manual_seed : Union[int, None] = None,
                     gpu_checkpoints: Union[list[str], None] = None,
+                    autocast : bool = False,
                     outsDataset: Union[dict[str, OutDatasetLoader], None] = {"default:Default" : OutDatasetLoader()},
                     images_log: list[str] = []) -> None:
         if os.environ["KONFAI_CONFIG_MODE"] != "Done":
@@ -328,6 +347,7 @@ class Predictor(DistributedObject):
         self.manual_seed = manual_seed
         self.dataset = dataset
         self.combine = combine
+        self.autocast = autocast
 
         self.model = model.getModel(train=False)
         self.it = 0
@@ -384,7 +404,7 @@ class Predictor(DistributedObject):
 
         shutil.copyfile(CONFIG_FILE(), self.predict_path+"Prediction.yml")
 
-        self.model.init(autocast=False, state = State.PREDICTION)
+        self.model.init(self.autocast, State.PREDICTION, self.dataset.getGroupsDest())
         self.model.init_outputsGroup()
         self.model._compute_channels_trace(self.model, self.model.in_channels, None, self.gpu_checkpoints)
         self.modelComposite = ModelComposite(self.model, len(MODEL().split(":")), self.combine)
@@ -402,7 +422,7 @@ class Predictor(DistributedObject):
     def run_process(self, world_size: int, global_rank: int, local_rank: int, dataloaders: list[DataLoader]):
         modelComposite = Network.to(self.modelComposite, local_rank*self.size)
         modelComposite = DDP(modelComposite, static_graph=True) if torch.cuda.is_available() else CPU_Model(modelComposite)
-        with _Predictor(world_size, global_rank, local_rank, self.predict_path, self.images_log, self.outsDataset, modelComposite, *dataloaders) as p:
+        with _Predictor(world_size, global_rank, local_rank, self.autocast, self.predict_path, self.images_log, self.outsDataset, modelComposite, *dataloaders) as p:
             p.run()
 
         

konfai/trainer.py

@@ -17,7 +17,6 @@ from konfai.utils.config import config
 from konfai.utils.utils import State, DataLog, DistributedObject, description, TrainerError
 from konfai.network.network import Network, ModelLoader, NetState, CPU_Model
 
-
 class EarlyStoppingBase:
 
     def __init__(self):
@@ -99,7 +98,7 @@ class _Trainer():
         self.autocast = autocast
         self.modelEMA = modelEMA
         self.early_stopping = EarlyStoppingBase() if early_stopping is None else early_stopping 
-
+        
         self.it_validation = it_validation
         if self.it_validation is None:
             self.it_validation = len(dataloader_training)
@@ -110,6 +109,8 @@ class _Trainer():
             for data in data_log:
                 self.data_log[data.split("/")[0].replace(":", ".")] = (DataLog.__getitem__(data.split("/")[1]).value[0], int(data.split("/")[2]))
         
+       
+        
     def __enter__(self):
         return self
     
@@ -118,13 +119,14 @@ class _Trainer():
             self.tb.close()
 
     def run(self) -> None:
-        with tqdm.tqdm(iterable = range(self.epoch, self.epochs), leave=False, total=self.epochs, initial=self.epoch, desc="Progress", disable=self.global_rank != 0) as epoch_tqdm:
+        self.dataloader_training.dataset.load("Train")
+        self.dataloader_validation.dataset.load("Validation")
+        with tqdm.tqdm(iterable = range(self.epoch, self.epochs), leave=False, total=self.epochs, initial=self.epoch, desc="Progress") as epoch_tqdm:
             for self.epoch in epoch_tqdm:
-                self.dataloader_training.dataset.load()
                 self.train()
                 if self.early_stopping.isStopped():
                     break
-                self.dataloader_training.dataset.resetAugmentation()
+                self.dataloader_training.dataset.resetAugmentation("Train")
                 
     def getInput(self, data_dict : dict[str, tuple[torch.Tensor, int, int, int, str, bool]]) -> dict[tuple[str, bool], torch.Tensor]:
         return {(k, v[5][0].item()) : v[0] for k, v in data_dict.items()}
@@ -137,7 +139,8 @@ class _Trainer():
             self.modelEMA.module.setState(NetState.TRAIN)
 
         desc = lambda : "Training : {}".format(description(self.model, self.modelEMA))
-        with tqdm.tqdm(iterable = enumerate(self.dataloader_training), desc = desc(), total=len(self.dataloader_training), leave=False, disable=self.global_rank != 0 and "KONFAI_CLUSTER" not in os.environ) as batch_iter:
+
+        with tqdm.tqdm(iterable = enumerate(self.dataloader_training), desc = desc(), total=len(self.dataloader_training), leave=False, ncols=0) as batch_iter:
             for _, data_dict in batch_iter:
                 with torch.amp.autocast('cuda', enabled=self.autocast):
                     input = self.getInput(data_dict)
@@ -171,8 +174,7 @@ class _Trainer():
 
         desc = lambda : "Validation : {}".format(description(self.model, self.modelEMA))
         data_dict = None
-        self.dataloader_validation.dataset.load()
-        with tqdm.tqdm(iterable = enumerate(self.dataloader_validation), desc = desc(), total=len(self.dataloader_validation), leave=False, disable=self.global_rank != 0 and "KONFAI_CLUSTER" not in os.environ) as batch_iter:
+        with tqdm.tqdm(iterable = enumerate(self.dataloader_validation), desc = desc(), total=len(self.dataloader_validation), leave=False, ncols=0) as batch_iter:
             for _, data_dict in batch_iter:
                 input = self.getInput(data_dict)
                 self.model(input)
@@ -180,7 +182,7 @@ class _Trainer():
                     self.modelEMA.module(input)
 
                 batch_iter.set_description(desc())
-        self.dataloader_validation.dataset.resetAugmentation()
+        self.dataloader_validation.dataset.resetAugmentation("Validation")
         dist.barrier()
         self.model.train()
         self.model.module.setState(NetState.TRAIN)
@@ -314,6 +316,18 @@ class Trainer(DistributedObject):
         self.ema_decay = ema_decay
         self.modelEMA : Union[torch.optim.swa_utils.AveragedModel, None] = None
         self.data_log = data_log
+
+        modules = []
+        for i,_ in self.model.named_modules():
+            modules.append(i)
+        for k in self.data_log:
+            tmp = k.split("/")[0].replace(":", ".")
+            if tmp not in self.dataset.getGroupsDest() and tmp not in modules:
+                raise TrainerError( f"Invalid key '{tmp}' in `data_log`.",
+                                   f"This key is neither a destination group from the dataset ({self.dataset.getGroupsDest()})",
+                                    f"nor a valid module name in the model ({modules}).",
+                "Please check your `data_log` configuration — it should reference either a model output or a dataset group.")
+            
         self.gradient_checkpoints = gradient_checkpoints
         self.gpu_checkpoints = gpu_checkpoints
         self.save_checkpoint_mode = save_checkpoint_mode
@@ -360,7 +374,7 @@ class Trainer(DistributedObject):
                     os.makedirs(dir)
 
             for name in sorted(os.listdir(path_checkpoint)):
-                checkpoint = torch.load(path_checkpoint+name, weights_only=False)
+                checkpoint = torch.load(path_checkpoint+name, weights_only=False, map_location='cpu')
                 self.model.load(checkpoint, init=False, ema=False)
 
                 torch.save(self.model, "{}Serialized/{}".format(path_model, name))

konfai/utils/dataset.py

@@ -348,6 +348,10 @@ class Dataset():
         @abstractmethod
         def getNames(self, group: str) -> list[str]:
             pass
+        
+        @abstractmethod
+        def getGroup(self) -> list[str]:
+            pass
 
         @abstractmethod
         def isExist(self, group: str, name: Union[str, None] = None) -> bool:
@@ -458,6 +462,9 @@ class Dataset():
                     names.extend(self.getNames("/".join(groups.split("/")[1:]), h5_group[group]))
             return names
         
+        def getGroup(self):
+            return self.h5.keys()
+        
         def _getDataset(self, groups: str, name: str, h5_group: h5py.Group = None) -> h5py.Dataset:
             if h5_group is None:
                 h5_group = self.h5
@@ -613,7 +620,10 @@ class Dataset():
         
         def getNames(self, group: str) -> list[str]:
             raise NotImplementedError()
-
+        
+        def getGroup(self):
+            raise NotImplementedError()
+        
         def getInfos(self, group: str, name: str) -> tuple[list[int], Attribute]:
             attributes = Attribute()
             if os.path.exists("{}{}{}.{}".format(self.filename, group if group is not None else "", name, self.format)):
@@ -736,7 +746,7 @@ class Dataset():
                 subDirectories.extend(self._getSubDirectories("/".join(groups.split("/")[1:]), subDirectory))
         return subDirectories
 
-    def getNames(self, groups: str, index: Union[list[int], None] = None, subDirectory: str = "") -> list[str]:
+    def getNames(self, groups: str, index: Union[list[int], None] = None) -> list[str]:
         names = []
         if self.is_directory:
             for subDirectory in self._getSubDirectories(groups):
@@ -752,6 +762,21 @@ class Dataset():
                 names = file.getNames(groups)
         return [name for i, name in enumerate(sorted(names)) if index is None or i in index]
     
+    def getGroup(self):
+        if self.is_directory:
+            groups = set()
+            for root, _, files in os.walk(self.filename):
+                for file in files:
+                    path = os.path.relpath(os.path.join(root, file.split(".")[0]), self.filename)
+                    parts = path.split("/")
+                    if len(parts) >= 2:
+                        del parts[-2]
+                    groups.add("/".join(parts))
+        else:
+            with Dataset.File(self.filename, True, self.format) as file:
+                groups = file.getGroup()
+        return list(groups)
+    
     def getInfos(self, groups: str, name: str) -> tuple[list[int], Attribute]:
         if self.is_directory:
             for subDirectory in self._getSubDirectories(groups):