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

reflectorch/models/networks/residual_net.py

@@ -1,157 +1,157 @@
-import torch
-from torch import nn
-from torch.nn import functional as F
-from torch.nn import init
-
-from reflectorch.models.activations import activation_by_name
-
-class ResidualMLP(nn.Module):
-    """Multilayer perceptron with residual blocks (BN-Act-Linear-BN-Act-Linear)"""
-
-    def __init__(
-            self,
-            dim_in: int,
-            dim_out: int,
-            dim_condition: int = 0,
-            layer_width: int = 512,
-            num_blocks: int = 4,
-            repeats_per_block: int = 2,
-            activation: str = 'relu',
-            use_batch_norm: bool = True,
-            use_layer_norm: bool = False,
-            dropout_rate: float = 0.0,
-            residual: bool = True,
-            adaptive_activation: bool = False,
-            conditioning: str = 'glu',
-            concat_condition_first_layer: bool = True,
-            film_with_tanh: bool = False,
-    ):
-        super().__init__()
-
-        self.concat_condition_first_layer = concat_condition_first_layer
-
-        dim_first_layer = dim_in + dim_condition if concat_condition_first_layer else dim_in
-        self.first_layer = nn.Linear(dim_first_layer, layer_width)
-
-        self.blocks = nn.ModuleList(
-            [
-                ResidualBlock(
-                    layer_width=layer_width,
-                    dim_condition=dim_condition,
-                    repeats_per_block=repeats_per_block,
-                    activation=activation,
-                    use_batch_norm=use_batch_norm,
-                    use_layer_norm=use_layer_norm,
-                    dropout_rate=dropout_rate,
-                    residual=residual,
-                    adaptive_activation=adaptive_activation,
-                    conditioning = conditioning,
-                    film_with_tanh = film_with_tanh,
-                )
-                for _ in range(num_blocks)
-            ]
-        )
-
-        self.last_layer = nn.Linear(layer_width, dim_out)
-
-    def forward(self, x, condition=None):
-        if self.concat_condition_first_layer and condition is not None:
-            x = self.first_layer(torch.cat([x, condition], dim=-1))
-        else:
-            x = self.first_layer(x)
-
-        for block in self.blocks:
-            x = block(x, condition=condition)
-
-        x = self.last_layer(x)
-
-        return x
-
-
-class ResidualBlock(nn.Module):
-    """Residual block (BN-Act-Linear-BN-Act-Linear)"""
-
-    def __init__(
-            self,
-            layer_width: int,
-            dim_condition: int = 0,
-            repeats_per_block: int = 2,
-            activation: str = 'relu',
-            use_batch_norm: bool = False,
-            use_layer_norm: bool = False,
-            dropout_rate: float = 0.0,
-            residual: bool = True,
-            adaptive_activation: bool = False,
-            conditioning: str = 'glu',
-            film_with_tanh: bool = False,
-    ):
-        super().__init__()
-         
-        self.residual = residual
-        self.repeats_per_block = repeats_per_block
-        self.use_batch_norm = use_batch_norm
-        self.use_layer_norm = use_layer_norm
-        self.dropout_rate = dropout_rate
-        self.adaptive_activation = adaptive_activation
-        self.conditioning = conditioning
-        self.film_with_tanh = film_with_tanh
-
-        if not adaptive_activation:
-            self.activation = activation_by_name(activation)()
-        else:
-            self.activation_layers = nn.ModuleList(
-                [activation_by_name(activation)() for _ in range(repeats_per_block)]
-            )
-
-        if use_batch_norm:
-            self.batch_norm_layers = nn.ModuleList(
-                [nn.BatchNorm1d(layer_width, eps=1e-3) for _ in range(repeats_per_block)]
-             )
-        elif use_layer_norm:
-            self.layer_norm_layers = nn.ModuleList(
-                [nn.LayerNorm(layer_width) for _ in range(repeats_per_block)]
-             )
-
-        if dim_condition:
-            if conditioning == 'glu':
-                self.condition_layer = nn.Linear(dim_condition, layer_width)
-            elif conditioning == 'film':
-                self.condition_layer = nn.Linear(dim_condition, 2*layer_width)
-
-        self.linear_layers = nn.ModuleList(
-            [nn.Linear(layer_width, layer_width) for _ in range(repeats_per_block)]
-        )
-
-        if self.dropout_rate > 0:
-            self.dropout = nn.Dropout(p=dropout_rate)
-
-    def forward(self, x, condition=None):
-        x0 = x
-         
-        for i in range(self.repeats_per_block):
-            if self.use_batch_norm:
-                x = self.batch_norm_layers[i](x)
-            elif self.use_layer_norm:
-                x = self.layer_norm_layers[i](x)
-
-            if not self.adaptive_activation:
-                x = self.activation(x)
-            else:
-                x = self.activation_layers[i](x)
-
-            if self.dropout_rate > 0 and i == self.repeats_per_block - 1:
-                x = self.dropout(x)
-
-            x = self.linear_layers[i](x)
-        
-        if condition is not None:
-            if self.conditioning == 'glu':
-                x = F.glu(torch.cat((x, self.condition_layer(condition)), dim=-1), dim=-1)
-            elif self.conditioning == 'film':
-                gamma, beta = torch.chunk(self.condition_layer(condition), chunks=2, dim=-1)
-                if self.film_with_tanh:
-                    tanh = nn.Tanh()
-                    gamma, beta = tanh(gamma), tanh(beta)
-                x = x * gamma + beta
-
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn import init
+
+from reflectorch.models.activations import activation_by_name
+
+class ResidualMLP(nn.Module):
+    """Multilayer perceptron with residual blocks (BN-Act-Linear-BN-Act-Linear)"""
+
+    def __init__(
+            self,
+            dim_in: int,
+            dim_out: int,
+            dim_condition: int = 0,
+            layer_width: int = 512,
+            num_blocks: int = 4,
+            repeats_per_block: int = 2,
+            activation: str = 'relu',
+            use_batch_norm: bool = True,
+            use_layer_norm: bool = False,
+            dropout_rate: float = 0.0,
+            residual: bool = True,
+            adaptive_activation: bool = False,
+            conditioning: str = 'glu',
+            concat_condition_first_layer: bool = True,
+            film_with_tanh: bool = False,
+    ):
+        super().__init__()
+
+        self.concat_condition_first_layer = concat_condition_first_layer
+
+        dim_first_layer = dim_in + dim_condition if concat_condition_first_layer else dim_in
+        self.first_layer = nn.Linear(dim_first_layer, layer_width)
+
+        self.blocks = nn.ModuleList(
+            [
+                ResidualBlock(
+                    layer_width=layer_width,
+                    dim_condition=dim_condition,
+                    repeats_per_block=repeats_per_block,
+                    activation=activation,
+                    use_batch_norm=use_batch_norm,
+                    use_layer_norm=use_layer_norm,
+                    dropout_rate=dropout_rate,
+                    residual=residual,
+                    adaptive_activation=adaptive_activation,
+                    conditioning = conditioning,
+                    film_with_tanh = film_with_tanh,
+                )
+                for _ in range(num_blocks)
+            ]
+        )
+
+        self.last_layer = nn.Linear(layer_width, dim_out)
+
+    def forward(self, x, condition=None):
+        if self.concat_condition_first_layer and condition is not None:
+            x = self.first_layer(torch.cat([x, condition], dim=-1))
+        else:
+            x = self.first_layer(x)
+
+        for block in self.blocks:
+            x = block(x, condition=condition)
+
+        x = self.last_layer(x)
+
+        return x
+
+
+class ResidualBlock(nn.Module):
+    """Residual block (BN-Act-Linear-BN-Act-Linear)"""
+
+    def __init__(
+            self,
+            layer_width: int,
+            dim_condition: int = 0,
+            repeats_per_block: int = 2,
+            activation: str = 'relu',
+            use_batch_norm: bool = False,
+            use_layer_norm: bool = False,
+            dropout_rate: float = 0.0,
+            residual: bool = True,
+            adaptive_activation: bool = False,
+            conditioning: str = 'glu',
+            film_with_tanh: bool = False,
+    ):
+        super().__init__()
+         
+        self.residual = residual
+        self.repeats_per_block = repeats_per_block
+        self.use_batch_norm = use_batch_norm
+        self.use_layer_norm = use_layer_norm
+        self.dropout_rate = dropout_rate
+        self.adaptive_activation = adaptive_activation
+        self.conditioning = conditioning
+        self.film_with_tanh = film_with_tanh
+
+        if not adaptive_activation:
+            self.activation = activation_by_name(activation)()
+        else:
+            self.activation_layers = nn.ModuleList(
+                [activation_by_name(activation)() for _ in range(repeats_per_block)]
+            )
+
+        if use_batch_norm:
+            self.batch_norm_layers = nn.ModuleList(
+                [nn.BatchNorm1d(layer_width, eps=1e-3) for _ in range(repeats_per_block)]
+             )
+        elif use_layer_norm:
+            self.layer_norm_layers = nn.ModuleList(
+                [nn.LayerNorm(layer_width) for _ in range(repeats_per_block)]
+             )
+
+        if dim_condition:
+            if conditioning == 'glu':
+                self.condition_layer = nn.Linear(dim_condition, layer_width)
+            elif conditioning == 'film':
+                self.condition_layer = nn.Linear(dim_condition, 2*layer_width)
+
+        self.linear_layers = nn.ModuleList(
+            [nn.Linear(layer_width, layer_width) for _ in range(repeats_per_block)]
+        )
+
+        if self.dropout_rate > 0:
+            self.dropout = nn.Dropout(p=dropout_rate)
+
+    def forward(self, x, condition=None):
+        x0 = x
+         
+        for i in range(self.repeats_per_block):
+            if self.use_batch_norm:
+                x = self.batch_norm_layers[i](x)
+            elif self.use_layer_norm:
+                x = self.layer_norm_layers[i](x)
+
+            if not self.adaptive_activation:
+                x = self.activation(x)
+            else:
+                x = self.activation_layers[i](x)
+
+            if self.dropout_rate > 0 and i == self.repeats_per_block - 1:
+                x = self.dropout(x)
+
+            x = self.linear_layers[i](x)
+        
+        if condition is not None:
+            if self.conditioning == 'glu':
+                x = F.glu(torch.cat((x, self.condition_layer(condition)), dim=-1), dim=-1)
+            elif self.conditioning == 'film':
+                gamma, beta = torch.chunk(self.condition_layer(condition), chunks=2, dim=-1)
+                if self.film_with_tanh:
+                    tanh = nn.Tanh()
+                    gamma, beta = tanh(gamma), tanh(beta)
+                x = x * gamma + beta
+
         return x0 + x if self.residual else x

reflectorch/paths.py

@@ -1,27 +1,29 @@
-from typing import Union
-from pathlib import Path
-
-__all__ = [
-    'ROOT_DIR',
-    'SAVED_MODELS_DIR',
-    'SAVED_LOSSES_DIR',
-    'RUN_SCRIPTS_DIR',
-    'CONFIG_DIR',
-    'TESTS_PATH',
-    'TEST_DATA_PATH',
-    'listdir',
-]
-
-ROOT_DIR: Path = Path(__file__).parents[1]
-SAVED_MODELS_DIR: Path = ROOT_DIR / 'saved_models'
-SAVED_LOSSES_DIR: Path = ROOT_DIR / 'saved_losses'
-RUN_SCRIPTS_DIR: Path = ROOT_DIR / 'runs'
-CONFIG_DIR: Path = ROOT_DIR / 'configs'
-TESTS_PATH: Path = ROOT_DIR / 'tests'
-TEST_DATA_PATH = TESTS_PATH / 'data'
-
-
-def listdir(path: Union[Path, str], pattern: str = '*', recursive: bool = False, *, sort_key=None, reverse=False):
-    path = Path(path)
-    func = path.rglob if recursive else path.glob
-    return sorted(list(func(pattern)), key=sort_key, reverse=reverse)
+from typing import Union
+from pathlib import Path
+
+__all__ = [
+    'ROOT_DIR',
+    'EXP_DATA_DIR',
+    'SAVED_MODELS_DIR',
+    'SAVED_LOSSES_DIR',
+    'RUN_SCRIPTS_DIR',
+    'CONFIG_DIR',
+    'TESTS_PATH',
+    'TEST_DATA_PATH',
+    'listdir',
+]
+
+ROOT_DIR: Path = Path(__file__).parents[1]
+EXP_DATA_DIR: Path = ROOT_DIR / 'exp_data'
+SAVED_MODELS_DIR: Path = ROOT_DIR / 'saved_models'
+SAVED_LOSSES_DIR: Path = ROOT_DIR / 'saved_losses'
+RUN_SCRIPTS_DIR: Path = ROOT_DIR / 'runs'
+CONFIG_DIR: Path = ROOT_DIR / 'configs'
+TESTS_PATH: Path = ROOT_DIR / 'tests'
+TEST_DATA_PATH = TESTS_PATH / 'data'
+
+
+def listdir(path: Union[Path, str], pattern: str = '*', recursive: bool = False, *, sort_key=None, reverse=False):
+    path = Path(path)
+    func = path.rglob if recursive else path.glob
+    return sorted(list(func(pattern)), key=sort_key, reverse=reverse)

reflectorch/runs/__init__.py

@@ -1,31 +1,31 @@
-from reflectorch.runs.train import (
-    run_train,
-    run_train_on_cluster,
-    run_test_config,
-)
-
-from reflectorch.runs.utils import (
-    train_from_config,
-    get_trainer_from_config,
-    get_paths_from_config,
-    get_callbacks_from_config,
-    get_trainer_by_name,
-    get_callbacks_by_name,
-    convert_pt_to_safetensors,
-)
-
-from reflectorch.runs.config import load_config
-
-__all__ = [
-    'run_train',
-    'run_train_on_cluster',
-    'train_from_config',
-    'run_test_config',
-    'get_trainer_from_config',
-    'get_paths_from_config',
-    'get_callbacks_from_config',
-    'get_trainer_by_name',
-    'get_callbacks_by_name',
-    'convert_pt_to_safetensors',
-    'load_config',
-]
+from reflectorch.runs.train import (
+    run_train,
+    run_train_on_cluster,
+    run_test_config,
+)
+
+from reflectorch.runs.utils import (
+    train_from_config,
+    get_trainer_from_config,
+    get_paths_from_config,
+    get_callbacks_from_config,
+    get_trainer_by_name,
+    get_callbacks_by_name,
+    convert_pt_to_safetensors,
+)
+
+from reflectorch.runs.config import load_config
+
+__all__ = [
+    'run_train',
+    'run_train_on_cluster',
+    'train_from_config',
+    'run_test_config',
+    'get_trainer_from_config',
+    'get_paths_from_config',
+    'get_callbacks_from_config',
+    'get_trainer_by_name',
+    'get_callbacks_by_name',
+    'convert_pt_to_safetensors',
+    'load_config',
+]

reflectorch/runs/config.py

@@ -1,25 +1,25 @@
-import yaml
-
-from pathlib import Path
-from reflectorch.paths import CONFIG_DIR
-
-
-def load_config(config_name: str, config_dir: str = None) -> dict:
-    """Loads a configuration dictionary from a YAML configuration file located in the configuration directory
-
-    Args:
-        config_name (str): name of the YAML configuration file
-        config_dir (str): path of the configuration directory
-
-    Returns:
-        dict: the configuration dictionary
-    """
-    if not config_name.endswith('.yaml'):
-        config_name = f'{config_name}.yaml'
-    config_dir = Path(config_dir) if config_dir else CONFIG_DIR
-    path = config_dir / config_name
-    with open(path, 'r') as f:
-        config = yaml.safe_load(f)
-    config['config_path'] = str(path.absolute())
-
-    return config
+import yaml
+
+from pathlib import Path
+from reflectorch.paths import CONFIG_DIR
+
+
+def load_config(config_name: str, config_dir: str = None) -> dict:
+    """Loads a configuration dictionary from a YAML configuration file located in the configuration directory
+
+    Args:
+        config_name (str): name of the YAML configuration file
+        config_dir (str): path of the configuration directory
+
+    Returns:
+        dict: the configuration dictionary
+    """
+    if not config_name.endswith('.yaml'):
+        config_name = f'{config_name}.yaml'
+    config_dir = Path(config_dir) if config_dir else CONFIG_DIR
+    path = config_dir / config_name
+    with open(path, 'r') as f:
+        config = yaml.safe_load(f)
+    config['config_path'] = str(path.absolute())
+
+    return config

reflectorch/runs/slurm_utils.py

@@ -1,93 +1,93 @@
-from typing import Tuple, Union
-from pathlib import Path
-import subprocess
-
-from reflectorch.paths import RUN_SCRIPTS_DIR
-
-
-def save_sbatch_and_run(
-        name: str,
-        args: str,
-        time: str,
-        partition: str = None,
-        reservation: bool = False,
-        chdir: str = '~/maxwell_output',
-        run_dir: Path = None,
-        confirm: bool = False,
-) -> Union[Tuple[str, str], None]:
-    run_dir = Path(run_dir) if run_dir else RUN_SCRIPTS_DIR
-    sbatch_path = run_dir / f'{name}.sh'
-
-    if sbatch_path.is_file():
-        import warnings
-        warnings.warn(f'Sbatch file {str(sbatch_path)} already exists!')
-        if confirm and not confirm_input('Continue?'):
-            return
-
-    file_content = _generate_sbatch_str(
-        name,
-        args,
-        time=time,
-        reservation=reservation,
-        partition=partition,
-        chdir=chdir,
-    )
-
-    if confirm:
-        print(f'Generated file content: \n{file_content}\n')
-        if not confirm_input(f'Save to {str(sbatch_path)} and run?'):
-            return
-
-    with open(str(sbatch_path), 'w') as f:
-        f.write(file_content)
-
-    res = submit_job(str(sbatch_path))
-    return res
-
-
-def _generate_sbatch_str(name: str,
-                         args: str,
-                         time: str,
-                         partition: str = None,
-                         reservation: bool = False,
-                         chdir: str = '~/maxwell_output',
-                         entry_point: str = 'python -m reflectorch.train',
-                         ) -> str:
-    chdir = str(Path(chdir).expanduser().absolute())
-    partition_keyword = 'reservation' if reservation else 'partition'
-
-    return f'''#!/bin/bash
-#SBATCH --chdir {chdir}
-#SBATCH --{partition_keyword}={partition}
-#SBATCH --constraint=P100
-#SBATCH --nodes=1
-#SBATCH --job-name {name}
-#SBATCH --time={time}
-#SBATCH --output {name}.out
-#SBATCH --error {name}.err
-
-{entry_point} {args}
-'''
-
-
-def confirm_input(message: str) -> bool:
-    yes = ('y', 'yes')
-    no = ('n', 'no')
-    res = ''
-    valid_results = list(yes) + list(no)
-    message = f'{message} Y/n: '
-
-    while res not in valid_results:
-        res = input(message).lower()
-    return res in yes
-
-
-def submit_job(sbatch_path: str) -> Tuple[str, str]:
-    process = subprocess.Popen(
-        ['sbatch', str(sbatch_path)],
-        stdout=subprocess.PIPE,
-        stderr=subprocess.PIPE,
-    )
-
-    stdout, stderr = process.communicate()
-    return stdout.decode(), stderr.decode()
+from typing import Tuple, Union
+from pathlib import Path
+import subprocess
+
+from reflectorch.paths import RUN_SCRIPTS_DIR
+
+
+def save_sbatch_and_run(
+        name: str,
+        args: str,
+        time: str,
+        partition: str = None,
+        reservation: bool = False,
+        chdir: str = '~/maxwell_output',
+        run_dir: Path = None,
+        confirm: bool = False,
+) -> Union[Tuple[str, str], None]:
+    run_dir = Path(run_dir) if run_dir else RUN_SCRIPTS_DIR
+    sbatch_path = run_dir / f'{name}.sh'
+
+    if sbatch_path.is_file():
+        import warnings
+        warnings.warn(f'Sbatch file {str(sbatch_path)} already exists!')
+        if confirm and not confirm_input('Continue?'):
+            return
+
+    file_content = _generate_sbatch_str(
+        name,
+        args,
+        time=time,
+        reservation=reservation,
+        partition=partition,
+        chdir=chdir,
+    )
+
+    if confirm:
+        print(f'Generated file content: \n{file_content}\n')
+        if not confirm_input(f'Save to {str(sbatch_path)} and run?'):
+            return
+
+    with open(str(sbatch_path), 'w') as f:
+        f.write(file_content)
+
+    res = submit_job(str(sbatch_path))
+    return res
+
+
+def _generate_sbatch_str(name: str,
+                         args: str,
+                         time: str,
+                         partition: str = None,
+                         reservation: bool = False,
+                         chdir: str = '~/maxwell_output',
+                         entry_point: str = 'python -m reflectorch.train',
+                         ) -> str:
+    chdir = str(Path(chdir).expanduser().absolute())
+    partition_keyword = 'reservation' if reservation else 'partition'
+
+    return f'''#!/bin/bash
+#SBATCH --chdir {chdir}
+#SBATCH --{partition_keyword}={partition}
+#SBATCH --constraint=P100
+#SBATCH --nodes=1
+#SBATCH --job-name {name}
+#SBATCH --time={time}
+#SBATCH --output {name}.out
+#SBATCH --error {name}.err
+
+{entry_point} {args}
+'''
+
+
+def confirm_input(message: str) -> bool:
+    yes = ('y', 'yes')
+    no = ('n', 'no')
+    res = ''
+    valid_results = list(yes) + list(no)
+    message = f'{message} Y/n: '
+
+    while res not in valid_results:
+        res = input(message).lower()
+    return res in yes
+
+
+def submit_job(sbatch_path: str) -> Tuple[str, str]:
+    process = subprocess.Popen(
+        ['sbatch', str(sbatch_path)],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+
+    stdout, stderr = process.communicate()
+    return stdout.decode(), stderr.decode()