Messy but working branch used for training models in parameter scan

README.md

@@ -15,6 +15,13 @@ conda install -c conda-forge tqdm
 conda install scipy
 ```
 
+To use `spline' flows, the third-party package *torchsearchsorted* must be installed.
+Download the source from the [github repo](https://github.com/aliutkus/torchsearchsorted) and, with the conda environment activated, run
+```bash
+pip install .
+```
+in the root folder.
+
 These are the minimal requirements for running the code, however if you plan
 on developing the code or doing small external tests, then I highly recommend also
 installing the following packages

anvil/benchmark_config/free_scalar_train.yml

@@ -8,29 +8,14 @@ m_sq: 4
 lam: 0
 use_arxiv_version: false
 
-# Base
-base: standard_normal
-
 # Model
+base: standard_normal
 model: real_nvp
-n_affine: 4
-
-# Networks
-standardise_inputs: false
-s_network_spec:
-    hidden_shape: [24]
-    activation: leaky_relu
-    final_activation: leaky_relu
-    do_batch_norm: false
-
-t_network_spec:
-    hidden_shape: [24]
-    activation: leaky_relu
-    final_activation: null
-    do_batch_norm: false
+model_spec:
+    n_affine: 2
 
 # Training length
-n_batch: 2000
+n_batch: 1000
 epochs: 3000
 save_interval: 1000
 

anvil/checkpoint.py

@@ -5,52 +5,16 @@
 is made so that we don't get unexpected results
 
 """
+from pathlib import Path
+from glob import glob
+
+import torch
+
+from reportengine.compat import yaml
 from copy import deepcopy
 
 from reportengine import collect
 
-from anvil.models import NeuralNetwork
-
-def neural_network(
-    size_half,
-    i_affine,
-    hidden_shape=(24,),
-    activation="leaky_relu",
-    final_activation=None,
-    do_batch_norm=False,
-):
-    """Returns an instance of NeuralNetwork to be used in real NVP
-
-    Parameters
-    ----------
-    size_half: int
-        Number of nodes in the input and output layer of the network
-    hidden_shape: list like
-        List like specifying the number of nodes in the intermediate layers
-    activation: (str, None)
-        Key representing the activation function used for each layer
-        except the final one.
-    final_activation: (str, None)
-        Key representing the activation function used on the final
-        layer.
-    do_batch_norm: bool
-        Flag dictating whether batch normalisation should be performed
-        before the activation function.
-    name: str
-        A label for the neural network, used for diagnostics.
-        """
-    return NeuralNetwork(
-        size_in=size_half,
-        size_out=size_half,
-        hidden_shape=hidden_shape,
-        activation=activation,
-        final_activation=final_activation,
-        do_batch_norm=do_batch_norm,
-        name=f"s{i_affine}"
-    )
-
-s_networks = collect("neural_network", ("affine_layer_index", "s_network_spec",))
-t_networks = collect("neural_network", ("affine_layer_index", "t_network_spec",))
 
 def loaded_checkpoint(checkpoint):
     if checkpoint is None:
@@ -62,14 +26,14 @@ def loaded_checkpoint(checkpoint):
 def train_range(loaded_checkpoint, epochs):
     if loaded_checkpoint is not None:
         cp_epoch = loaded_checkpoint["epoch"]
-        train_range = (cp_epoch, cp_epoch + epochs)
+        train_range = (cp_epoch, epochs)
     else:
         train_range = (0, epochs)
     return train_range
 
 
-def loaded_model(loaded_checkpoint, flow_model):
-    new_model = deepcopy(flow_model)  # need to copy model so we don't get weird results
+def loaded_model(loaded_checkpoint, model_to_load):
+    new_model = deepcopy(model_to_load)  # need to copy model so we don't get weird results
     if loaded_checkpoint is not None:
         new_model.load_state_dict(loaded_checkpoint["model_state_dict"])
     return new_model
@@ -79,3 +43,76 @@ def current_loss(loaded_checkpoint):
     if loaded_checkpoint is None:
         return None
     return loaded_checkpoint["loss"]
+
+
+class InvalidCheckpointError(Exception):
+    pass
+
+
+class InvalidTrainingOutputError(Exception):
+    pass
+
+
+class TrainingRuncardNotFound(InvalidTrainingOutputError):
+    pass
+
+
+class Checkpoint:
+    """Class which saves and loads checkpoints and allows checkpoints to be
+    sorted"""
+
+    def __init__(self, path: str):
+        self.path = Path(path)
+        try:
+            self.epoch = int(self.path.stem.split("_")[-1])  # should be an int
+        except ValueError:
+            raise InvalidCheckpointError(
+                f"{self.path} does not match expected "
+                "name checkpoint: `checkpoint_<epoch>.pt`"
+            )
+
+    def __lt__(self, other):
+        return self.epoch < other.epoch
+
+    def __repr__(self):
+        return str(self.path)
+
+    def load(self):
+        """Return checkpoint dictionary"""
+        return torch.load(self.path)
+
+
+class TrainingOutput:
+    """Class which acts as container for training output, which is a directory
+    containing training configuration, checkpoints and training logs
+    """
+
+    _loaded_config = None
+
+    def __init__(self, path: str):
+        self.path = Path(path)
+        self.config = self.path / "runcard.yml"
+        if not self.config.is_file():
+            raise TrainingRuncardNotFound(
+                f"Invalid training output, no runcard found at: {self.config}"
+            )
+        self.checkpoints = [
+            Checkpoint(cp_path) for cp_path in glob(f"{self.path}/checkpoints/*")
+        ]
+        self.cp_ids = [cp.epoch for cp in self.checkpoints]
+        self.name = self.path.name
+
+    def get_config(self):
+        if self._loaded_config is None:
+            with open(self.config, "r") as f:
+                self._loaded_config = yaml.safe_load(f)
+        return self._loaded_config
+
+    def as_input(self):
+        inp = dict(self.get_config())  # make copy
+        inp["checkpoints"] = self.checkpoints
+        inp["cp_ids"] = self.cp_ids
+        return inp
+
+    def final_checkpoint(self):
+        return max(self.checkpoints)

anvil/config.py

@@ -8,11 +8,13 @@
 from reportengine.report import Config
 from reportengine.configparser import ConfigError, element_of, explicit_node
 
-from anvil.core import TrainingOutput
-from anvil.train import OPTIMIZER_OPTIONS, reduce_lr_on_plateau
-from anvil.models import MODEL_OPTIONS
+from anvil.core import normalising_flow
 from anvil.geometry import Geometry2D
+from anvil.checkpoint import TrainingOutput
+from anvil.train import OPTIMIZER_OPTIONS, SCHEDULER_OPTIONS
+from anvil.models import MODEL_OPTIONS
 from anvil.distributions import BASE_OPTIONS, TARGET_OPTIONS
+from anvil.fields import FIELD_OPTIONS
 
 log = logging.getLogger(__name__)
 
@@ -36,7 +38,7 @@ def produce_lattice_size(self, lattice_length, lattice_dimension):
         return pow(lattice_length, lattice_dimension)
 
     def produce_config_size(self, lattice_size, target):
-        """number of nodes in a single field configuration"""
+        """Size of a single configuration or input vector for neural network."""
         if target == "o3":
             return 2 * lattice_size
         return lattice_size
@@ -53,16 +55,8 @@ def produce_size_half(self, config_size):
     def produce_geometry(self, lattice_length):
         return Geometry2D(lattice_length)
 
-    def parse_target(self, target: str):
-        """String specifying target distrbution."""
-        return target
-
-    def parse_base(self, base: str):
-        """String specifying base distribution."""
-        return base
-
     @explicit_node
-    def produce_target_dist(self, target: str):
+    def produce_target_dist(self, target):
         """Return the function which initialises the correct action"""
         try:
             return TARGET_OPTIONS[target]
@@ -71,6 +65,18 @@ def produce_target_dist(self, target: str):
                 f"invalid target distribution {target}", target, TARGET_OPTIONS.keys()
             )
 
+    @explicit_node
+    def produce_field(self, target):
+        """Return the function which instantiates the field object, for
+        calculating observables."""
+        try:
+            return FIELD_OPTIONS[target]
+        except KeyError:
+            log.warning(
+                f"Target {target} does not match an implemented field theory. Using generic field class."
+            )
+            return FIELD_OPTIONS[None]
+
     @explicit_node
     def produce_base_dist(self, base: str):
         """Return the action which loads appropriate base distribution"""
@@ -101,54 +107,26 @@ def parse_radius(self, rad: (int, float, str)):
         """Radius for semicircle distribution."""
         return rad
 
-    def parse_m_sq(self, m: (float, int)):
-        """Bare mass squared in scalar theory."""
-        return m
-
-    def parse_lam(self, lam: (float, int)):
-        """Coefficient of quartic interaction in phi^4 theory."""
-        return lam
-
-    def parse_use_arxiv_version(self, do_use: bool):
-        """If true, use the conventional phi^4 action. If false,
-        there is an additional factor of 1/2 for the kinetic part
-        of the phi^4 action."""
-        return do_use
-
-    def parse_beta(self, beta: (float, int)):
-        """Inverse temperature."""
-        return beta
+    def parse_couplings(self, couplings: dict):
+        """Couplings for field theory."""
+        return couplings  # TODO: obviously need to be more fool-proof about this
 
-    def parse_model(self, model: str):
-        """Label for normalising flow model."""
-        return model
+    def parse_parameterisation(self, param: str):
+        return param
 
     @explicit_node
-    def produce_flow_model(self, model):
-        """Return the action which instantiates the normalising flow model."""
+    def produce_model_action(self, model: str):
+        """Given a string, return the flow model action indexed by that string."""
         try:
             return MODEL_OPTIONS[model]
         except KeyError:
-            raise ConfigError(
-                f"invalid flow model {model}", model, MODEL_OPTIONS.keys()
-            )
-
-    def parse_standardise_inputs(self, do_stand: bool):
-        """Flag specifying whether to standardise input vectors before
-        passing them through a neural network."""
-        return do_stand
+            raise ConfigError(f"Invalid model {model}", model, MODEL_OPTIONS.keys())
 
-    def parse_n_affine(self, n: int):
-        """Number of affine layers."""
-        return n
-
-    def produce_affine_layer_index(self, n_affine):
-        """Given n_affine, the number of affine layers, produces a list
-        with n_affine elements, the ith element is {i_affine: i}
-
-        we can use affine_layer_index to collect over when producing the model
-        """
-        return [{"i_affine": i} for i in range(n_affine)]
+    @explicit_node
+    def produce_model_to_load(self, n_mixture=1):
+        """Produce the generative model, whose parameters are to be loaded, which maps
+        the base to an approximate of the target distribution."""
+        return normalising_flow
 
     def parse_n_batch(self, nb: int):
         """Batch size for training."""
@@ -206,9 +184,14 @@ def produce_loaded_optimizer(self, optimizer):
             )
 
     @explicit_node
-    def produce_scheduler(self):
+    def produce_loaded_scheduler(self, scheduler):
         """Currently fixed to ReduceLROnPlateau"""
-        return reduce_lr_on_plateau
+        try:
+            return SCHEDULER_OPTIONS[scheduler]
+        except KeyError:
+            raise ConfigError(
+                f"Invalid scheduler {scheduler}", scheduler, SCHEDULER_OPTIONS.keys()
+            )
 
     def parse_target_length(self, targ: int):
         """Target number of decorrelated field configurations to generate."""
@@ -248,6 +231,9 @@ def parse_n_boot(self, n_boot: int):
         log.warning(f"Using user specified n_boot: {n_boot}")
         return n_boot
 
+    def parse_connected_correlator(self, connected: bool):
+        return connected
+
     @element_of("windows")
     def parse_window(self, window: float):
         """A numerical factor featuring in the calculation of the optimal 'window'