tensorneko 0.3.5__tar.gz → 0.3.6__tar.gz

{tensorneko-0.3.5/src/tensorneko.egg-info → tensorneko-0.3.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tensorneko
-Version: 0.3.5
+Version: 0.3.6
 Summary: Tensor Neural Engine Kompanion. An util library based on PyTorch and PyTorch Lightning.
 Home-page: https://github.com/ControlNet/tensorneko
 Author: ControlNet
@@ -38,7 +38,7 @@ License-File: LICENSE
 <div align="center">    
     <a href="https://www.python.org/"><img src="https://img.shields.io/pypi/pyversions/tensorneko?style=flat-square"></a>
     <a href="https://pytorch.org/"><img src="https://img.shields.io/badge/PyTorch-%3E%3D1.9.0-EE4C2C?style=flat-square&logo=pytorch"></a>
-    <a href="https://www.pytorchlightning.ai/"><img src="https://img.shields.io/badge/Lightning-2.0.*/2.1.*-792EE5?style=flat-square&logo=lightning"></a>
+    <a href="https://www.pytorchlightning.ai/"><img src="https://img.shields.io/badge/Lightning-2.0.*%20|%202.1.*-792EE5?style=flat-square&logo=lightning"></a>
 </div>
 
 <div align="center">

{tensorneko-0.3.5 → tensorneko-0.3.6}/README.md

@@ -12,7 +12,7 @@
 <div align="center">    
     <a href="https://www.python.org/"><img src="https://img.shields.io/pypi/pyversions/tensorneko?style=flat-square"></a>
     <a href="https://pytorch.org/"><img src="https://img.shields.io/badge/PyTorch-%3E%3D1.9.0-EE4C2C?style=flat-square&logo=pytorch"></a>
-    <a href="https://www.pytorchlightning.ai/"><img src="https://img.shields.io/badge/Lightning-2.0.*/2.1.*-792EE5?style=flat-square&logo=lightning"></a>
+    <a href="https://www.pytorchlightning.ai/"><img src="https://img.shields.io/badge/Lightning-2.0.*%20|%202.1.*-792EE5?style=flat-square&logo=lightning"></a>
 </div>
 
 <div align="center">

{tensorneko-0.3.5 → tensorneko-0.3.6}/src/tensorneko/arch/binary_classifier.py

@@ -1,18 +1,17 @@
-from abc import ABC
-from typing import Optional, Union, Sequence, Dict
+from typing import Optional, Union, Sequence, Dict, Any
 
 from torch import Tensor
-from torch.nn import BCEWithLogitsLoss
+from torch.nn import BCEWithLogitsLoss, Module
 from torch.optim import Adam
 from torchmetrics import Accuracy, F1Score, AUROC
 
 from ..neko_model import NekoModel
 
 
-class BinaryClassifier(NekoModel, ABC):
+class BinaryClassifier(NekoModel):
 
-    def __init__(self, model=None, learning_rate: float = 1e-4, distributed: bool = False):
-        super().__init__()
+    def __init__(self, name, model: Module, learning_rate: float = 1e-4, distributed: bool = False):
+        super().__init__(name)
         self.save_hyperparameters()
         self.model = model
         self.learning_rate = learning_rate
@@ -23,8 +22,10 @@ class BinaryClassifier(NekoModel, ABC):
         self.auc_fn = AUROC(task="binary")
 
     @classmethod
-    def from_module(cls, model, learning_rate: float = 1e-4, distributed=False):
-        return cls(model, learning_rate, distributed)
+    def from_module(cls, model: Module, learning_rate: float = 1e-4, name: str = "binary_classifier",
+        distributed: bool = False
+    ):
+        return cls(name, model, learning_rate, distributed)
 
     def forward(self, x):
         return self.model(x)
@@ -49,6 +50,10 @@ class BinaryClassifier(NekoModel, ABC):
     ) -> Dict[str, Tensor]:
         return self.step(batch)
 
+    def predict_step(self, batch: Tensor, batch_idx: int, dataloader_idx: Optional[int] = None) -> Any:
+        x, y = batch
+        return self(x)
+
     def configure_optimizers(self):
         optimizer = Adam(self.parameters(), lr=self.learning_rate)
         return [optimizer]

tensorneko-0.3.6/src/tensorneko/callback/gpu_stats_logger.py

@@ -0,0 +1,68 @@
+from typing import Any
+
+from lightning.pytorch import Callback, Trainer, LightningModule
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+
+
+class GpuStatsLogger(Callback):
+    """Log GPU stats for each training epoch"""
+
+    def __init__(self, delay: float = 0.5, on_epoch: bool = True, on_step: bool = False):
+        try:
+            from gpumonitor.monitor import GPUStatMonitor
+        except ImportError:
+            raise ImportError("gpumonitor is required to use GPUStatsLogger")
+
+        self.monitor_epoch = GPUStatMonitor(delay=delay) if on_epoch else None
+        self.monitor_step = GPUStatMonitor(delay=delay) if on_step else None
+        self.on_epoch = on_epoch
+        self.on_step = on_step
+        assert self.on_epoch or self.on_step, "on_epoch and on_step cannot be both False"
+
+    def on_train_epoch_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        if not self.on_epoch:
+            return
+        self.monitor_epoch.reset()
+
+    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        if not self.on_epoch:
+            return
+        for gpu in self.monitor_epoch.average_stats.gpus:
+            logged_info = {
+                f"gpu{gpu.index}_memory_used_epoch": gpu.memory_used / 1024,
+                f"gpu{gpu.index}_memory_total_epoch": gpu.memory_total / 1024,
+                f"gpu{gpu.index}_memory_util_epoch": gpu.memory_used / gpu.memory_total,
+                f"gpu{gpu.index}_temperature_epoch": float(gpu.temperature),
+                f"gpu{gpu.index}_utilization_epoch": gpu.utilization / 100,
+                f"gpu{gpu.index}_power_draw_epoch": float(gpu.power_draw),
+                f"gpu{gpu.index}_power_percentage_epoch": gpu.power_draw / gpu.power_limit,
+                f"gpu{gpu.index}_fan_speed_epoch": float(gpu.fan_speed) if gpu.fan_speed is not None else 0.,
+            }
+            pl_module.logger.log_metrics(logged_info, step=trainer.global_step)
+            pl_module.log_dict(logged_info, logger=False, sync_dist=pl_module.distributed)
+
+    def on_train_batch_start(
+        self, trainer: Trainer, pl_module: LightningModule, batch: Any, batch_idx: int
+    ) -> None:
+        if not self.on_step:
+            return
+        self.monitor_step.reset()
+
+    def on_train_batch_end(
+        self, trainer: Trainer, pl_module: LightningModule, outputs: STEP_OUTPUT, batch: Any, batch_idx: int
+    ) -> None:
+        if not self.on_step:
+            return
+        for gpu in self.monitor_step.average_stats.gpus:
+            logged_info = {
+                f"gpu{gpu.index}_memory_used_step": gpu.memory_used / 1024,
+                f"gpu{gpu.index}_memory_total_step": gpu.memory_total / 1024,
+                f"gpu{gpu.index}_memory_util_step": gpu.memory_used / gpu.memory_total,
+                f"gpu{gpu.index}_temperature_step": float(gpu.temperature),
+                f"gpu{gpu.index}_utilization_step": gpu.utilization / 100,
+                f"gpu{gpu.index}_power_draw_step": float(gpu.power_draw),
+                f"gpu{gpu.index}_power_percentage_step": gpu.power_draw / gpu.power_limit,
+                f"gpu{gpu.index}_fan_speed_step": float(gpu.fan_speed) if gpu.fan_speed is not None else 0.,
+            }
+            pl_module.logger.log_metrics(logged_info, step=trainer.global_step)
+            pl_module.log_dict(logged_info, logger=False, sync_dist=pl_module.distributed)

tensorneko-0.3.6/src/tensorneko/callback/system_stats_logger.py

@@ -0,0 +1,44 @@
+from typing import Any
+
+from lightning.pytorch import Callback, Trainer, LightningModule
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+
+
+class SystemStatsLogger(Callback):
+    """Log system stats for each training epoch"""
+
+    def __init__(self, on_epoch: bool = True, on_step: bool = False):
+        try:
+            import psutil
+        except ImportError:
+            raise ImportError("psutil is required to use SystemStatsLogger")
+        self.psutil = psutil
+        self.on_epoch = on_epoch
+        self.on_step = on_step
+        assert self.on_epoch or self.on_step, "on_epoch and on_step cannot be both False"
+
+    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        if not self.on_epoch:
+            return
+        cpu_usage = self.psutil.cpu_percent()
+        memory_usage = self.psutil.virtual_memory().percent
+        logged_info = {
+            "cpu_usage_epoch": cpu_usage,
+            "memory_usage_epoch": memory_usage
+        }
+        pl_module.logger.log_metrics(logged_info, step=trainer.global_step)
+        pl_module.log_dict(logged_info, logger=False, sync_dist=pl_module.distributed)
+
+    def on_train_batch_end(
+        self, trainer: Trainer, pl_module: LightningModule, outputs: STEP_OUTPUT, batch: Any, batch_idx: int
+    ) -> None:
+        if not self.on_step:
+            return
+        cpu_usage = self.psutil.cpu_percent()
+        memory_usage = self.psutil.virtual_memory().percent
+        logged_info = {
+            "cpu_usage_step": cpu_usage,
+            "memory_usage_step": memory_usage
+        }
+        pl_module.logger.log_metrics(logged_info, step=trainer.global_step)
+        pl_module.log_dict(logged_info, logger=False, sync_dist=pl_module.distributed)

{tensorneko-0.3.5 → tensorneko-0.3.6}/src/tensorneko/evaluation/fid.py

@@ -90,8 +90,8 @@ class FID:
 
         if progress_bar:
             tqdm = import_tqdm_auto().tqdm
-            pred = tqdm(total=len(pred), desc="Forward predicted features")
-            true = tqdm(total=len(true), desc="Forward ground truth features")
+            pred = tqdm(pred, total=len(pred), desc="Forward predicted features")
+            true = tqdm(true, total=len(true), desc="Forward ground truth features")
 
         for batch in pred:
             self.fid.update(batch.to(self.device), real=False)

{tensorneko-0.3.5 → tensorneko-0.3.6}/src/tensorneko/neko_trainer.py

@@ -3,16 +3,20 @@ from datetime import timedelta
 from time import time
 from typing import Optional, Union, List, Dict
 
+from lightning.fabric.plugins.precision.precision import _PRECISION_INPUT
+from lightning.fabric.utilities.types import _PATH
 from lightning.pytorch import Trainer, Callback
 from lightning.pytorch.accelerators import Accelerator
 from lightning.pytorch.callbacks import ModelCheckpoint, Checkpoint
 from lightning.pytorch.loggers import Logger, TensorBoardLogger
-from lightning.pytorch.plugins import PLUGIN_INPUT
 from lightning.pytorch.profilers import Profiler
 from lightning.pytorch.strategies import Strategy
 from lightning.pytorch.trainer.connectors.accelerator_connector import _LITERAL_WARN
-from lightning.fabric.plugins.precision.precision import _PRECISION_INPUT
-from lightning.fabric.utilities.types import _PATH
+
+try:
+    from lightning.pytorch.plugins import PLUGIN_INPUT
+except ImportError:
+    from lightning.pytorch.plugins import _PLUGIN_INPUT as PLUGIN_INPUT
 
 from .callback import NilCallback, LrLogger, EpochNumLogger, EpochTimeLogger, GpuStatsLogger, SystemStatsLogger
 

{tensorneko-0.3.5 → tensorneko-0.3.6}/src/tensorneko/util/__init__.py

@@ -9,7 +9,7 @@ from . import type
 from .configuration import Configuration
 from .misc import reduce_dict_by, summarize_dict_by, with_printed_shape, is_bad_num, count_parameters, compose, \
     generate_inf_seq, listdir, with_printed, ifelse, dict_add, as_list, identity, list_to_dict, circular_pad, \
-    load_py, try_until_success
+    load_py, try_until_success, sample_indexes
 from .misc import get_tensorneko_path
 from .dispatched_misc import sparse2binary, binary2sparse
 from .reproducibility import Seed
@@ -71,6 +71,7 @@ __all__ = [
     "circular_pad",
     "load_py",
     "try_until_success",
+    "sample_indexes",
     "download_file",
     "WindowMerger",
 ]

{tensorneko-0.3.5 → tensorneko-0.3.6}/src/tensorneko/util/configuration.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Any
+from typing import Generic
 
+from tensorneko_util.util.type import T
 
-class Configuration(ABC):
+
+class Configuration(ABC, Generic[T]):
     """
     Configuration base abstract class.
 
@@ -55,7 +57,7 @@ class Configuration(ABC):
         return iter((*self.args, *self.kwargs.values()))
 
     @abstractmethod
-    def build(self) -> Any:
+    def build(self) -> T:
         """
         A method to build an object.
 

{tensorneko-0.3.5 → tensorneko-0.3.6}/src/tensorneko/util/misc.py

@@ -9,7 +9,7 @@ from torch import Tensor
 from torch.nn import Module
 
 from tensorneko_util.util.misc import generate_inf_seq, listdir, with_printed, ifelse, dict_add, as_list, \
-    identity, list_to_dict, compose, circular_pad, load_py, try_until_success
+    identity, list_to_dict, compose, circular_pad, load_py, try_until_success, sample_indexes
 from .type import T, A
 
 
@@ -165,3 +165,4 @@ list_to_dict = list_to_dict
 circular_pad = circular_pad
 load_py = load_py
 try_until_success = try_until_success
+sample_indexes = sample_indexes

{tensorneko-0.3.5 → tensorneko-0.3.6}/src/tensorneko/util/type.py

@@ -1,4 +1,3 @@
-from enum import Enum
 from typing import Callable, Union, List, Tuple, TypeVar
 
 import numpy as np

tensorneko-0.3.6/src/tensorneko/version.txt

@@ -0,0 +1 @@
+0.3.6

{tensorneko-0.3.5 → tensorneko-0.3.6/src/tensorneko.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tensorneko
-Version: 0.3.5
+Version: 0.3.6
 Summary: Tensor Neural Engine Kompanion. An util library based on PyTorch and PyTorch Lightning.
 Home-page: https://github.com/ControlNet/tensorneko
 Author: ControlNet
@@ -38,7 +38,7 @@ License-File: LICENSE
 <div align="center">    
     <a href="https://www.python.org/"><img src="https://img.shields.io/pypi/pyversions/tensorneko?style=flat-square"></a>
     <a href="https://pytorch.org/"><img src="https://img.shields.io/badge/PyTorch-%3E%3D1.9.0-EE4C2C?style=flat-square&logo=pytorch"></a>
-    <a href="https://www.pytorchlightning.ai/"><img src="https://img.shields.io/badge/Lightning-2.0.*/2.1.*-792EE5?style=flat-square&logo=lightning"></a>
+    <a href="https://www.pytorchlightning.ai/"><img src="https://img.shields.io/badge/Lightning-2.0.*%20|%202.1.*-792EE5?style=flat-square&logo=lightning"></a>
 </div>
 
 <div align="center">

{tensorneko-0.3.5 → tensorneko-0.3.6}/src/tensorneko.egg-info/requires.txt

@@ -8,7 +8,7 @@ pillow>=8.1
 av>=8.0.3
 numpy>=1.20.1
 einops>=0.3.0
-tensorneko_util==0.3.5
+tensorneko_util==0.3.6
 
 [:platform_system == "Windows"]
 pysoundfile>=0.9.0

tensorneko-0.3.5/src/tensorneko/callback/gpu_stats_logger.py

@@ -1,31 +0,0 @@
-from lightning.pytorch import Callback, Trainer, LightningModule
-
-
-class GpuStatsLogger(Callback):
-    """Log GPU stats for each training epoch"""
-
-    def __init__(self, delay: float = 0.5):
-        try:
-            from gpumonitor.monitor import GPUStatMonitor
-        except ImportError:
-            raise ImportError("gpumonitor is required to use GPUStatsLogger")
-
-        self.monitor = GPUStatMonitor(delay=delay)
-
-    def on_train_epoch_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
-        self.monitor.reset()
-
-    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
-        for gpu in self.monitor.average_stats.gpus:
-            logged_info = {
-                f"gpu{gpu.index}_memory_used": gpu.memory_used / 1024,
-                f"gpu{gpu.index}_memory_total": gpu.memory_total / 1024,
-                f"gpu{gpu.index}_memory_util": gpu.memory_used / gpu.memory_total,
-                f"gpu{gpu.index}_temperature": float(gpu.temperature),
-                f"gpu{gpu.index}_utilization": gpu.utilization / 100,
-                f"gpu{gpu.index}_power_draw": float(gpu.power_draw),
-                f"gpu{gpu.index}_power_percentage": gpu.power_draw / gpu.power_limit,
-                f"gpu{gpu.index}_fan_speed": float(gpu.fan_speed) if gpu.fan_speed is not None else 0.,
-            }
-            pl_module.logger.log_metrics(logged_info, step=trainer.global_step)
-            pl_module.log_dict(logged_info, logger=False, sync_dist=pl_module.distributed)

tensorneko-0.3.5/src/tensorneko/callback/system_stats_logger.py

@@ -1,22 +0,0 @@
-from lightning.pytorch import Callback, Trainer, LightningModule
-
-
-class SystemStatsLogger(Callback):
-    """Log system stats for each training epoch"""
-
-    def __init__(self):
-        try:
-            import psutil
-        except ImportError:
-            raise ImportError("psutil is required to use SystemStatsLogger")
-        self.psutil = psutil
-
-    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
-        cpu_usage = self.psutil.cpu_percent()
-        memory_usage = self.psutil.virtual_memory().percent
-        logged_info = {
-            "cpu_usage": cpu_usage,
-            "memory_usage": memory_usage
-        }
-        pl_module.logger.log_metrics(logged_info, step=trainer.global_step)
-        pl_module.log_dict(logged_info, logger=False, sync_dist=pl_module.distributed)

tensorneko-0.3.5/src/tensorneko/version.txt

@@ -1 +0,0 @@
-0.3.5