reflectorch 1.3.0__py3-none-any.whl → 1.5.0__py3-none-any.whl

reflectorch/ml/callbacks.py

@@ -1,81 +1,81 @@
-import torch
-
-import numpy as np
-
-from reflectorch.ml.basic_trainer import (
-    TrainerCallback,
-    Trainer,
-)
-from reflectorch.ml.utils import is_divisor
-
-__all__ = [
-    'SaveBestModel',
-    'LogLosses',
-]
-
-
-class SaveBestModel(TrainerCallback):
-    """Callback for periodically saving the best model weights
-
-    Args:
-        path (str): path for saving the model weights
-        freq (int, optional): frequency in iterations at which the current average loss is evaluated. Defaults to 50.
-        average (int, optional): number of recent iterations over which the average loss is computed. Defaults to 10.
-    """
-
-    def __init__(self, path: str, freq: int = 50, average: int = 10):
-        self.path = path
-        self.average = average
-        self._best_loss = np.inf
-        self.freq = freq
-
-    def end_batch(self, trainer: Trainer, batch_num: int) -> None:
-        """checks if the current average loss has improved from the previous save, if true the model is saved
-
-        Args:
-            trainer (Trainer): the trainer object
-            batch_num (int): the current iteration / batch
-        """
-        if is_divisor(batch_num, self.freq):
-
-            loss = np.mean(trainer.losses['total_loss'][-self.average:])
-
-            if loss < self._best_loss:
-                self._best_loss = loss
-                self.save(trainer, batch_num)
-
-    def save(self, trainer: Trainer, batch_num: int):
-        """saves a dictionary containing the network weights, the learning rates, the losses and the current \
-            best loss with its corresponding iteration to the disk
-
-        Args:
-            trainer (Trainer): the trainer object
-            batch_num (int): the current iteration / batch
-        """
-        prev_save = trainer.callback_params.pop('saved_iteration', 0)
-        trainer.callback_params['saved_iteration'] = batch_num
-        save_dict = {
-            'model': trainer.model.state_dict(),
-            'lrs': trainer.lrs,
-            'losses': trainer.losses,
-            'prev_save': prev_save,
-            'batch_num': batch_num,
-            'best_loss': self._best_loss
-        }
-        torch.save(save_dict, self.path)
-
-
-class LogLosses(TrainerCallback):
-    """Callback for logging the training losses"""
-    def end_batch(self, trainer: Trainer, batch_num: int) -> None:
-        """log loss at the current iteration
-
-        Args:
-            trainer (Trainer): the trainer object
-            batch_num (int): the index of the current iteration / batch
-        """
-        for loss_name, loss_values in trainer.losses.items():
-            try:
-                trainer.log(f'train/{loss_name}', loss_values[-1])
-            except IndexError:
+import torch
+
+import numpy as np
+
+from reflectorch.ml.basic_trainer import (
+    TrainerCallback,
+    Trainer,
+)
+from reflectorch.ml.utils import is_divisor
+
+__all__ = [
+    'SaveBestModel',
+    'LogLosses',
+]
+
+
+class SaveBestModel(TrainerCallback):
+    """Callback for periodically saving the best model weights
+
+    Args:
+        path (str): path for saving the model weights
+        freq (int, optional): frequency in iterations at which the current average loss is evaluated. Defaults to 50.
+        average (int, optional): number of recent iterations over which the average loss is computed. Defaults to 10.
+    """
+
+    def __init__(self, path: str, freq: int = 50, average: int = 10):
+        self.path = path
+        self.average = average
+        self._best_loss = np.inf
+        self.freq = freq
+
+    def end_batch(self, trainer: Trainer, batch_num: int) -> None:
+        """checks if the current average loss has improved from the previous save, if true the model is saved
+
+        Args:
+            trainer (Trainer): the trainer object
+            batch_num (int): the current iteration / batch
+        """
+        if is_divisor(batch_num, self.freq):
+
+            loss = np.mean(trainer.losses['total_loss'][-self.average:])
+
+            if loss < self._best_loss:
+                self._best_loss = loss
+                self.save(trainer, batch_num)
+
+    def save(self, trainer: Trainer, batch_num: int):
+        """saves a dictionary containing the network weights, the learning rates, the losses and the current \
+            best loss with its corresponding iteration to the disk
+
+        Args:
+            trainer (Trainer): the trainer object
+            batch_num (int): the current iteration / batch
+        """
+        prev_save = trainer.callback_params.pop('saved_iteration', 0)
+        trainer.callback_params['saved_iteration'] = batch_num
+        save_dict = {
+            'model': trainer.model.state_dict(),
+            'lrs': trainer.lrs,
+            'losses': trainer.losses,
+            'prev_save': prev_save,
+            'batch_num': batch_num,
+            'best_loss': self._best_loss
+        }
+        torch.save(save_dict, self.path)
+
+
+class LogLosses(TrainerCallback):
+    """Callback for logging the training losses"""
+    def end_batch(self, trainer: Trainer, batch_num: int) -> None:
+        """log loss at the current iteration
+
+        Args:
+            trainer (Trainer): the trainer object
+            batch_num (int): the index of the current iteration / batch
+        """
+        for loss_name, loss_values in trainer.losses.items():
+            try:
+                trainer.log(f'train/{loss_name}', loss_values[-1])
+            except IndexError:
                 continue

reflectorch/ml/dataloaders.py

@@ -1,27 +1,27 @@
-from torch import Tensor
-
-
-from reflectorch.data_generation import BasicDataset
-from reflectorch.data_generation.reflectivity import kinematical_approximation
-from reflectorch.data_generation.priors import BasicParams
-from reflectorch.ml.basic_trainer import DataLoader
-
-
-__all__ = [
-    "ReflectivityDataLoader",
-    "MultilayerDataLoader",
-]
-
-
-class ReflectivityDataLoader(BasicDataset, DataLoader):
-    """Dataloader for reflectivity data, combining functionality from the ``BasicDataset`` (basic dataset class for reflectivity) and the ``DataLoader`` (which inherits from ``TrainerCallback``) classes"""
-    pass
-
-
-class MultilayerDataLoader(ReflectivityDataLoader):
-    """Dataloader for reflectivity curves simulated using the kinematical approximation"""
-    def _sample_from_prior(self, batch_size: int):
-        return self.prior_sampler.optimized_sample(batch_size)
-
-    def _calc_curves(self, q_values: Tensor, params: BasicParams):
+from torch import Tensor
+
+
+from reflectorch.data_generation import BasicDataset
+from reflectorch.data_generation.reflectivity import kinematical_approximation
+from reflectorch.data_generation.priors import BasicParams
+from reflectorch.ml.basic_trainer import DataLoader
+
+
+__all__ = [
+    "ReflectivityDataLoader",
+    "MultilayerDataLoader",
+]
+
+
+class ReflectivityDataLoader(BasicDataset, DataLoader):
+    """Dataloader for reflectivity data, combining functionality from the ``BasicDataset`` (basic dataset class for reflectivity) and the ``DataLoader`` (which inherits from ``TrainerCallback``) classes"""
+    pass
+
+
+class MultilayerDataLoader(ReflectivityDataLoader):
+    """Dataloader for reflectivity curves simulated using the kinematical approximation"""
+    def _sample_from_prior(self, batch_size: int):
+        return self.prior_sampler.optimized_sample(batch_size)
+
+    def _calc_curves(self, q_values: Tensor, params: BasicParams):
         return kinematical_approximation(q_values, params.thicknesses, params.roughnesses, params.slds)

reflectorch/ml/loggers.py

@@ -1,56 +1,56 @@
-from torch.utils.tensorboard import SummaryWriter
-
-__all__ = [
-    'Logger',
-    'Loggers',
-    'PrintLogger',
-    'TensorBoardLogger',
-]
-
-
-class Logger(object):
-    "Base class defining a common interface for logging"
-    def log(self, name: str, data):
-        pass
-
-    def __setitem__(self, key, value):
-        """Enable dictionary-style setting to log data."""
-        self.log(key, value)
-
-
-class Loggers(Logger):
-    """Class for using multiple loggers"""
-    def __init__(self, *loggers):
-        self._loggers = tuple(loggers)
-
-    def log(self, name: str, data):
-        for logger in self._loggers:
-            logger.log(name, data)
-
-
-class PrintLogger(Logger):
-    """Logger which prints to the console"""
-    def log(self, name: str, data):
-        print(name, ': ', data)
-
-class TensorBoardLogger(Logger):
-    def __init__(self, log_dir: str):
-        """
-        Args:
-            log_dir (str): Directory where TensorBoard logs will be written
-        """
-        super().__init__()
-        self.writer = SummaryWriter(log_dir=log_dir)
-        self.step = 1
-        
-    def log(self, name: str, data):
-        """Log scalar data to TensorBoard
-        
-        Args:
-            name (str): Name/tag for the data
-            data: Scalar value to log
-        """
-        if hasattr(data, 'item'):
-            data = data.item()
-        self.writer.add_scalar(name, data, self.step)
+from torch.utils.tensorboard import SummaryWriter
+
+__all__ = [
+    'Logger',
+    'Loggers',
+    'PrintLogger',
+    'TensorBoardLogger',
+]
+
+
+class Logger(object):
+    "Base class defining a common interface for logging"
+    def log(self, name: str, data):
+        pass
+
+    def __setitem__(self, key, value):
+        """Enable dictionary-style setting to log data."""
+        self.log(key, value)
+
+
+class Loggers(Logger):
+    """Class for using multiple loggers"""
+    def __init__(self, *loggers):
+        self._loggers = tuple(loggers)
+
+    def log(self, name: str, data):
+        for logger in self._loggers:
+            logger.log(name, data)
+
+
+class PrintLogger(Logger):
+    """Logger which prints to the console"""
+    def log(self, name: str, data):
+        print(name, ': ', data)
+
+class TensorBoardLogger(Logger):
+    def __init__(self, log_dir: str):
+        """
+        Args:
+            log_dir (str): Directory where TensorBoard logs will be written
+        """
+        super().__init__()
+        self.writer = SummaryWriter(log_dir=log_dir)
+        self.step = 1
+        
+    def log(self, name: str, data):
+        """Log scalar data to TensorBoard
+        
+        Args:
+            name (str): Name/tag for the data
+            data: Scalar value to log
+        """
+        if hasattr(data, 'item'):
+            data = data.item()
+        self.writer.add_scalar(name, data, self.step)
         self.step += 1