tensorneko 0.3.7__tar.gz → 0.3.9__tar.gz

{tensorneko-0.3.7/src/tensorneko.egg-info → tensorneko-0.3.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tensorneko
-Version: 0.3.7
+Version: 0.3.9
 Summary: Tensor Neural Engine Kompanion. An util library based on PyTorch and PyTorch Lightning.
 Home-page: https://github.com/ControlNet/tensorneko
 Author: ControlNet

{tensorneko-0.3.7 → tensorneko-0.3.9}/src/tensorneko/__init__.py

@@ -1,11 +1,11 @@
 import os.path
 
-from tensorneko_util import io
 from . import backend
 from . import callback
 from . import dataset
 from . import debug
 from . import evaluation
+from . import io
 from . import layer
 from . import module
 from . import notebook
@@ -13,10 +13,13 @@ from . import optim
 from . import preprocess
 from . import util
 from . import visualization
+from .io import read, write
 from .neko_model import NekoModel
 from .neko_module import NekoModule
 from .neko_trainer import NekoTrainer
 
+__version__ = io.read.text(os.path.join(util.get_tensorneko_path(), "version.txt"))
+
 __all__ = [
     "callback",
     "dataset",
@@ -33,7 +36,9 @@ __all__ = [
     "debug",
     "NekoModel",
     "NekoTrainer",
-    "NekoModule"
+    "NekoModule",
+    "read",
+    "write",
 ]
 
-__version__ = io.read.text(os.path.join(util.get_tensorneko_path(), "version.txt"))
+

{tensorneko-0.3.7 → tensorneko-0.3.9}/src/tensorneko/arch/binary_classifier.py

@@ -6,16 +6,18 @@ from torch.optim import Adam
 from torchmetrics import Accuracy, F1Score, AUROC
 
 from ..neko_model import NekoModel
+from ..util import Shape
 
 
 class BinaryClassifier(NekoModel):
 
-    def __init__(self, name, model: Module, learning_rate: float = 1e-4, distributed: bool = False):
-        super().__init__(name)
+    def __init__(self, name, model: Module, learning_rate: float = 1e-4, distributed: bool = False,
+        input_shape: Optional[Shape] = None
+    ):
+        super().__init__(name, input_shape=input_shape, distributed=distributed)
         self.save_hyperparameters()
         self.model = model
         self.learning_rate = learning_rate
-        self.distributed = distributed
         self.loss_fn = BCEWithLogitsLoss()
         self.acc_fn = Accuracy(task="binary")
         self.f1_fn = F1Score(task="binary")
@@ -23,9 +25,9 @@ class BinaryClassifier(NekoModel):
 
     @classmethod
     def from_module(cls, model: Module, learning_rate: float = 1e-4, name: str = "binary_classifier",
-        distributed: bool = False
+        distributed: bool = False, input_shape: Optional[Shape] = None
     ):
-        return cls(name, model, learning_rate, distributed)
+        return cls(name, model, learning_rate, distributed, input_shape)
 
     def forward(self, x):
         return self.model(x)
@@ -38,7 +40,7 @@ class BinaryClassifier(NekoModel):
         acc = self.acc_fn(prob, y)
         f1 = self.f1_fn(prob, y)
         auc = self.auc_fn(prob, y)
-        return {"loss": loss, "acc": acc, "f1": f1, "auc": auc}
+        return {"loss": loss, "metric/acc": acc, "metric/f1": f1, "metric/auc": auc}
 
     def training_step(self, batch: Optional[Union[Tensor, Sequence[Tensor]]] = None, batch_idx: Optional[int] = None,
         optimizer_idx: Optional[int] = None, hiddens: Optional[Tensor] = None

{tensorneko-0.3.7 → tensorneko-0.3.9}/src/tensorneko/arch/gan.py

@@ -140,7 +140,7 @@ class GAN(NekoModel, ABC):
 
         # calculate discriminator loss
         d_loss = self.d_loss_fn(self.discriminator(d_batch), d_label)
-        return {"loss": d_loss, "d_loss": d_loss}
+        return {"loss": d_loss, "loss/d_loss": d_loss}
 
     def g_step(self, x: Tensor, z: Tensor) -> Dict[str, Tensor]:
         # forward generator
@@ -150,7 +150,7 @@ class GAN(NekoModel, ABC):
 
         # calculate generator loss
         g_loss = self.g_loss_fn(self.discriminator(fake_image), target_label)
-        return {"loss": g_loss, "g_loss": g_loss}
+        return {"loss": g_loss, "loss/g_loss": g_loss}
 
     def validation_step(self, batch: Optional[Union[Tensor, Sequence[Tensor]]] = None, batch_idx: Optional[int] = None,
         dataloader_idx: Optional[int] = None

{tensorneko-0.3.7 → tensorneko-0.3.9}/src/tensorneko/arch/vqvae.py

@@ -72,7 +72,7 @@ class VQVAE(NekoModel, ABC):
         x_hat, embedding_loss = self(x)
         rec_loss = self.rec_loss_fn(x_hat, x)
         loss = rec_loss + embedding_loss
-        return {"loss": loss, "r_loss": rec_loss, "e_loss": embedding_loss}
+        return {"loss": loss, "loss/r_loss": rec_loss, "loss/e_loss": embedding_loss}
 
     def validation_step(self, batch: Optional[Union[Tensor, Sequence[Tensor]]] = None, batch_idx: Optional[int] = None,
         dataloader_idx: Optional[int] = None
@@ -84,7 +84,7 @@ class VQVAE(NekoModel, ABC):
         x_hat, embedding_loss = self(x)
         rec_loss = self.rec_loss_fn(x_hat, x)
         loss = rec_loss + embedding_loss
-        return {"loss": loss, "r_loss": rec_loss, "e_loss": embedding_loss}
+        return {"loss": loss, "loss/r_loss": rec_loss, "loss/e_loss": embedding_loss}
 
     def predict_step(self, batch: Tensor, batch_idx: int, dataloader_idx: Optional[int] = None) -> Tensor:
         return self(batch)[0]

{tensorneko-0.3.7 → tensorneko-0.3.9}/src/tensorneko/arch/wgan.py

@@ -100,4 +100,4 @@ class WGAN(GAN, ABC):
             gp = torch.tensor(0., device=self.device)
             loss = d_loss
 
-        return {"loss": loss, "d_loss": d_loss, "gp": gp}
+        return {"loss": loss, "loss/d_loss": d_loss, "loss/gp": gp}

{tensorneko-0.3.7 → tensorneko-0.3.9}/src/tensorneko/backend/__init__.py

@@ -1,8 +1,10 @@
-from tensorneko_util.backend import parallel, run_blocking, VisualLib, AudioLib
+from tensorneko_util.backend import parallel, run_blocking, VisualLib, AudioLib, import_tqdm_auto, import_tqdm
 
 __all__ = [
     "parallel",
     "run_blocking",
     "VisualLib",
     "AudioLib",
+    "import_tqdm_auto",
+    "import_tqdm",
 ]

{tensorneko-0.3.7 → tensorneko-0.3.9}/src/tensorneko/callback/earlystop_lr.py

@@ -22,7 +22,7 @@ class EarlyStoppingLR(Callback):
             return
         all_lr = []
         for key, value in metrics.items():
-            if re.match(r"opt\d+_lr\d+", key):
+            if re.match(r"learning_rate/opt\d+_lr\d+", key):
                 all_lr.append(value)
 
         if len(all_lr) == 0:

{tensorneko-0.3.7 → tensorneko-0.3.9}/src/tensorneko/callback/epoch_num_logger.py

@@ -5,7 +5,7 @@ class EpochNumLogger(Callback):
     """Log epoch number in each epoch start."""
 
     def on_train_epoch_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
-        key = "epoch"
+        key = "epoch/epoch"
         value = trainer.current_epoch
         pl_module.logger.log_metrics({key: value}, step=trainer.global_step)
         pl_module.log(key, value, logger=False, sync_dist=pl_module.distributed)

{tensorneko-0.3.7 → tensorneko-0.3.9}/src/tensorneko/callback/epoch_time_logger.py

@@ -12,5 +12,5 @@ class EpochTimeLogger(Callback):
 
     def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
         elapsed_time = time() - self.start_time
-        pl_module.logger.log_metrics({"epoch_time": elapsed_time}, step=trainer.global_step)
-        pl_module.log("epoch_time", elapsed_time, logger=False, sync_dist=pl_module.distributed)
+        pl_module.logger.log_metrics({"epoch/epoch_time": elapsed_time}, step=trainer.global_step)
+        pl_module.log("epoch/epoch_time", elapsed_time, logger=False, sync_dist=pl_module.distributed)

{tensorneko-0.3.7 → tensorneko-0.3.9}/src/tensorneko/callback/gpu_stats_logger.py

@@ -29,14 +29,14 @@ class GpuStatsLogger(Callback):
             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.,
+                f"system/gpu{gpu.index}_memory_used_epoch": gpu.memory_used / 1024,
+                f"system/gpu{gpu.index}_memory_total_epoch": gpu.memory_total / 1024,
+                f"system/gpu{gpu.index}_memory_util_epoch": gpu.memory_used / gpu.memory_total,
+                f"system/gpu{gpu.index}_temperature_epoch": float(gpu.temperature),
+                f"system/gpu{gpu.index}_utilization_epoch": gpu.utilization / 100,
+                f"system/gpu{gpu.index}_power_draw_epoch": float(gpu.power_draw),
+                f"system/gpu{gpu.index}_power_percentage_epoch": gpu.power_draw / gpu.power_limit,
+                f"system/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)
@@ -55,14 +55,14 @@ class GpuStatsLogger(Callback):
             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.,
+                f"system/gpu{gpu.index}_memory_used_step": gpu.memory_used / 1024,
+                f"system/gpu{gpu.index}_memory_total_step": gpu.memory_total / 1024,
+                f"system/gpu{gpu.index}_memory_util_step": gpu.memory_used / gpu.memory_total,
+                f"system/gpu{gpu.index}_temperature_step": float(gpu.temperature),
+                f"system/gpu{gpu.index}_utilization_step": gpu.utilization / 100,
+                f"system/gpu{gpu.index}_power_draw_step": float(gpu.power_draw),
+                f"system/gpu{gpu.index}_power_percentage_step": gpu.power_draw / gpu.power_limit,
+                f"system/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.7 → tensorneko-0.3.9}/src/tensorneko/callback/lr_logger.py

@@ -7,7 +7,7 @@ class LrLogger(Callback):
     def on_train_epoch_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
         for i, optimizer in enumerate(trainer.optimizers):
             for j, params in enumerate(optimizer.param_groups):
-                key = f"opt{i}_lr{j}"
+                key = f"learning_rate/opt{i}_lr{j}"
                 value = params["lr"]
                 pl_module.logger.log_metrics({key: value}, step=trainer.global_step)
                 pl_module.log(key, value, logger=False, sync_dist=pl_module.distributed)

{tensorneko-0.3.7 → tensorneko-0.3.9}/src/tensorneko/callback/system_stats_logger.py

@@ -23,8 +23,8 @@ class SystemStatsLogger(Callback):
         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
+            "system/cpu_usage_epoch": cpu_usage,
+            "system/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)
@@ -37,8 +37,8 @@ class SystemStatsLogger(Callback):
         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
+            "system/cpu_usage_step": cpu_usage,
+            "system/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.7 → tensorneko-0.3.9}/src/tensorneko/evaluation/fid.py

@@ -12,12 +12,13 @@ from torchmetrics.image.fid import FrechetInceptionDistance
 from torchvision.transforms.functional import resize
 
 from tensorneko_util.backend import VisualLib
-from tensorneko_util.backend._tqdm import import_tqdm_auto
+from tensorneko_util.backend.tqdm import import_tqdm_auto
 
 try:
     from typing import Literal
     TypeOption = Literal["video", "image"]
 except ImportError:
+    Literal = None
     TypeOption = str
 
 

{tensorneko-0.3.7 → tensorneko-0.3.9}/src/tensorneko/io/reader.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 from typing import Type
 
 from tensorneko_util.io.reader import Reader as BaseReader
+from .weight import WeightReader
 
 try:
     from .mesh import MeshReader
@@ -15,6 +16,7 @@ class Reader(BaseReader):
 
     def __init__(self):
         super().__init__()
+        self.weight = WeightReader
         self._mesh = None
 
     @property

tensorneko-0.3.9/src/tensorneko/io/weight/__init__.py

@@ -0,0 +1,2 @@
+from .weight_reader import WeightReader
+from .weight_writer import WeightWriter

tensorneko-0.3.9/src/tensorneko/io/weight/weight_reader.py

@@ -0,0 +1,81 @@
+from typing import OrderedDict
+
+import torch
+
+from ...util import Device
+
+
+class WeightReader:
+    """WeightReader for read model weights (checkpoints, state_dict, etc)."""
+
+    @classmethod
+    def of_pt(cls, path: str, map_location: Device = "cpu") -> OrderedDict[str, torch.Tensor]:
+        """
+        Reads PyTorch model weights from a `.pt` or `.pth` file.
+
+        Args:
+            path (``str``): The path of the `.pt` or `.pth` file.
+            map_location (:class:`torch.device` | ``str``): The location to load the model weights. Default: "cpu"
+
+        Returns:
+            :class:`collections.OrderedDict`[``str``, :class:`torch.Tensor`]: The model weights.
+        """
+        return torch.load(path, map_location=map_location)
+
+    @classmethod
+    def of_ckpt(cls, path: str, map_location: Device = "cpu") -> OrderedDict[str, torch.Tensor]:
+        """
+        Reads PyTorch model weights from a `.ckpt` file.
+
+        Args:
+            path (``str``): The path of the `.ckpt` file.
+            map_location (:class:`torch.device` | ``str``): The location to load the model weights. Default: "cpu"
+
+        Returns:
+            :class:`collections.OrderedDict`[``str``, :class:`torch.Tensor`]: The model weights.
+        """
+        return torch.load(path, map_location=map_location)["state_dict"]
+
+    @classmethod
+    def of_safetensors(cls, path: str, map_location: str = "cpu") -> OrderedDict[str, torch.Tensor]:
+        """
+        Reads model weights from a `.safetensors` file.
+
+        Args:
+            path (``str``): The path of the `.safetensors` file.
+            map_location (``str``): The location to load the model weights. Default: "cpu"
+
+        Returns:
+            :class:`collections.OrderedDict`[``str``, :class:`torch.Tensor`]: The model weights.
+        """
+        import safetensors
+        from collections import OrderedDict
+        tensors = OrderedDict()
+        with safetensors.safe_open(path, framework="pt", device=map_location) as f:
+            for key in f.keys():
+                tensors[key] = f.get_tensor(key)
+        return tensors
+
+    @classmethod
+    def of(cls, path: str, map_location: Device = "cpu") -> OrderedDict[str, torch.Tensor]:
+        """
+        Reads model weights from a file.
+
+        Args:
+            path (``str``): The path of the file.
+            map_location (:class:`torch.device` | ``str``): The location to load the model weights. Default: "cpu"
+
+        Returns:
+            :class:`collections.OrderedDict`[``str``, :class:`torch.Tensor`]: The model weights.
+        """
+
+        if path.endswith(".pt") or path.endswith(".pth"):
+            return cls.of_pt(path, map_location)
+        elif path.endswith(".ckpt"):
+            return cls.of_ckpt(path, map_location)
+        elif path.endswith(".safetensors"):
+            if isinstance(map_location, torch.device):
+                map_location = str(map_location)
+            return cls.of_safetensors(path, map_location)
+        else:
+            raise ValueError("Unknown file type. Supported types: .pt, .pth, .ckpt, .safetensors")

tensorneko-0.3.9/src/tensorneko/io/weight/weight_writer.py

@@ -0,0 +1,48 @@
+from typing import Dict
+
+import torch
+
+
+class WeightWriter:
+    """WeightWriter for write model weights (checkpoints, state_dict, etc)."""
+
+    @classmethod
+    def to_pt(cls, path: str, weights: Dict[str, torch.Tensor]) -> None:
+        """
+        Writes PyTorch model weights to a `.pt` or `.pth` file.
+
+        Args:
+            path (``str``): The path of the `.pt` or `.pth` file.
+            weights (:class:`collections.OrderedDict`[``str``, :class:`torch.Tensor`]): The model weights.
+        """
+        torch.save(weights, path)
+
+    @classmethod
+    def to_safetensors(cls, path: str, weights: Dict[str, torch.Tensor]) -> None:
+        """
+        Writes model weights to a `.safetensors` file.
+
+        Args:
+            path (``str``): The path of the `.safetensors` file.
+            weights (:class:`collections.OrderedDict`[``str``, :class:`torch.Tensor`]): The model weights.
+        """
+        import safetensors.torch
+        safetensors.torch.save_file(weights, path)
+
+    @classmethod
+    def to(cls, path: str, weights: Dict[str, torch.Tensor]) -> None:
+        """
+        Writes model weights to a file.
+
+        Args:
+            path (``str``): The path of the file.
+            weights (:class:`collections.OrderedDict`[``str``, :class:`torch.Tensor`]): The model weights.
+        """
+        file_type = path.split(".")[-1]
+
+        if file_type == "pt":
+            cls.to_pt(path, weights)
+        elif file_type == "safetensors":
+            cls.to_safetensors(path, weights)
+        else:
+            raise ValueError("Unknown file type. Supported types: .pt, .safetensors")

{tensorneko-0.3.7 → tensorneko-0.3.9}/src/tensorneko/io/writer.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from typing import Type
 
 from tensorneko_util.io.writer import Writer as BaseWriter
+from .weight import WeightWriter
 
 try:
     from .mesh import MeshWriter
@@ -16,6 +17,7 @@ class Writer(BaseWriter):
 
     def __init__(self):
         super().__init__()
+        self.weight = WeightWriter
         self._mesh = None
 
     @property

{tensorneko-0.3.7 → tensorneko-0.3.9}/src/tensorneko/layer/noise.py

@@ -22,7 +22,7 @@ class GaussianNoise(NekoModule):
         from https://discuss.pytorch.org/t/writing-a-simple-gaussian-noise-layer-in-pytorch/4694/3
     """
 
-    def __init__(self, sigma=0.1, device: Union[Device, str] = "cuda"):
+    def __init__(self, sigma=0.1, device: Device = "cuda"):
         super().__init__()
         self.sigma = sigma
         self.noise = torch.tensor(0.).to(device)

{tensorneko-0.3.7 → tensorneko-0.3.9}/src/tensorneko/neko_model.py

@@ -170,6 +170,7 @@ class NekoModel(LightningModule, NekoModule):
             and self.logger is not None \
             and self.trainer.global_step % self.trainer.log_every_n_steps == 0:
             self.log_on_training_step_end(outputs)
+        super().on_train_batch_end(outputs, batch, batch_idx)
 
     def on_validation_batch_end(
         self, outputs: Optional[STEP_OUTPUT], batch: Any, batch_idx: int, dataloader_idx: int = 0
@@ -177,18 +178,21 @@ class NekoModel(LightningModule, NekoModule):
         """Append the outputs of validation step to the list"""
         outputs = {k: v.detach() for k, v in outputs.items()}
         self.validation_step_outputs.append(outputs)
+        super().on_validation_batch_end(outputs, batch, batch_idx, dataloader_idx)
 
     def on_train_epoch_end(self) -> None:
         """For each training epoch end, log the metrics"""
         if self.trainer.log_on_epoch and self.logger is not None:
             self.log_on_training_epoch_end(self.training_step_outputs)
         self.training_step_outputs.clear()
+        super().on_train_epoch_end()
 
     def on_validation_epoch_end(self) -> None:
         """For each validation epoch end, log the metrics"""
         if self.logger is not None:
             self.log_on_validation_epoch_end(self.validation_step_outputs)
         self.validation_step_outputs.clear()
+        super().on_validation_epoch_end()
 
     def log_on_training_epoch_end(self, outputs: EPOCH_OUTPUT) -> None:
         """Log the training epoch outputs"""
@@ -197,7 +201,8 @@ class NekoModel(LightningModule, NekoModule):
             getter = summarize_dict_by(key, torch.mean)
             value = getter(outputs)
             history_item[key] = value
-            self.logger.log_metrics({key: value}, step=self.trainer.global_step)
+            if self.logger is not None:
+                self.logger.log_metrics({key: value}, step=self.trainer.global_step)
             self.log(key, value, on_epoch=True, on_step=False, logger=False, sync_dist=self.distributed)
         self.history.append(history_item)
 
@@ -210,7 +215,8 @@ class NekoModel(LightningModule, NekoModule):
             getter = summarize_dict_by(key, torch.mean)
             value = getter(outputs)
             self.history[-1]["val_" + key] = value
-            self.logger.log_metrics({key: value}, step=self.trainer.global_step)
+            if self.logger is not None:
+                self.logger.log_metrics({key: value}, step=self.trainer.global_step)
             self.log(key, value, on_epoch=True, on_step=False, logger=False, sync_dist=self.distributed)
             self.log(f"val_{key}", value, on_epoch=True, on_step=False, logger=False, prog_bar=True,
                 sync_dist=self.distributed)
@@ -220,7 +226,8 @@ class NekoModel(LightningModule, NekoModule):
         history_item = {}
         for key, value in output.items():
             history_item[key] = value
-            self.logger.log_metrics({key: value}, step=self.trainer.global_step)
+            if self.logger is not None:
+                self.logger.log_metrics({key: value}, step=self.trainer.global_step)
             self.log(key, value, on_epoch=False, on_step=True, logger=False, prog_bar=key == "loss",
                 sync_dist=self.distributed)
         self.history.append(history_item)
@@ -230,6 +237,7 @@ class NekoModel(LightningModule, NekoModule):
     ) -> None:
         """For each test step end, log the metrics"""
         self.log_dict(outputs, sync_dist=self.distributed)
+        super().on_test_batch_end(outputs, batch, batch_idx, dataloader_idx)
 
     def log_image(self, name: str, image: torch.Tensor) -> None:
         """Log an image to the logger"""
@@ -238,6 +246,15 @@ class NekoModel(LightningModule, NekoModule):
 
         self.logger.experiment.add_image(name, torch.clip(image, 0, 1), self.trainer.global_step)
 
+    def log_histogram(self, name: str, values: torch.Tensor) -> None:
+        """Log a histogram to the logger"""
+        if self.logger is None:
+            return
+
+        self.logger.experiment.add_histogram(name, values, self.trainer.global_step)
+
     def on_train_start(self) -> None:
         """Log the model graph to tensorboard when the input shape is set"""
-        self.logger.log_graph(self, self.example_input_array)
+        if self.can_plot_graph and self.logger is not None:
+            self.logger.log_graph(self, self.example_input_array)
+        super().on_train_start()

{tensorneko-0.3.7 → tensorneko-0.3.9}/src/tensorneko/util/type.py

@@ -11,7 +11,7 @@ ModuleFactory = Union[Callable[[], Module], Callable[[int], Module]]
 """The module builder type of ``() -> torch.nn.Module | (int) -> torch.nn.Module``"""
 
 
-Device = device
+Device = Union[str, device]
 """Device type of :class:`torch.device`"""
 
 

tensorneko-0.3.9/src/tensorneko/version.txt

@@ -0,0 +1 @@
+0.3.9

{tensorneko-0.3.7 → tensorneko-0.3.9/src/tensorneko.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tensorneko
-Version: 0.3.7
+Version: 0.3.9
 Summary: Tensor Neural Engine Kompanion. An util library based on PyTorch and PyTorch Lightning.
 Home-page: https://github.com/ControlNet/tensorneko
 Author: ControlNet

{tensorneko-0.3.7 → tensorneko-0.3.9}/src/tensorneko.egg-info/SOURCES.txt

@@ -47,6 +47,9 @@ src/tensorneko/io/writer.py
 src/tensorneko/io/mesh/__init__.py
 src/tensorneko/io/mesh/mesh_reader.py
 src/tensorneko/io/mesh/mesh_writer.py
+src/tensorneko/io/weight/__init__.py
+src/tensorneko/io/weight/weight_reader.py
+src/tensorneko/io/weight/weight_writer.py
 src/tensorneko/layer/__init__.py
 src/tensorneko/layer/aggregation.py
 src/tensorneko/layer/attention.py

{tensorneko-0.3.7 → tensorneko-0.3.9}/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.7
+tensorneko_util==0.3.9
 
 [:platform_system == "Windows"]
 pysoundfile>=0.9.0

tensorneko-0.3.7/src/tensorneko/version.txt

@@ -1 +0,0 @@
-0.3.7