Training

Solvers for training, testing, checkpointing and recovering of networks.

MultiTaskSolver

Implements training, testing, checkpoint, recovering of flyvis networks.

Gives access to the network, decoder, task, optimizer, penalty and scheduler and the directory where the results are stored.

Parameters:

Name	Type	Description	Default
name	str	Name of the solver.	''
config	Optional[Union[dict, Namespace]]	Configuration for the solver.	None
init_network	bool	Whether to initialize the network. Defaults to True.	True
init_decoder	bool	Whether to initialize the decoder. Defaults to True.	True
init_task	bool	Whether to initialize the task. Defaults to True.	True
init_optim	bool	Whether to initialize the optimizer. Defaults to True.	True
init_penalties	bool	Whether to initialize penalties. Defaults to True.	True
init_scheduler	bool	Whether to initialize the scheduler. Defaults to True.	True
delete_if_exists	bool	Whether to delete existing directory. Defaults to False.	False

Attributes:

Name	Type	Description
dir	NetworkDir	Directory where results are stored.
network	Network	The neural network.
decoder	Dict[str, Module]	The decoder modules.
task	Task	The task being solved.
optimizer	Optimizer	The optimizer.
penalty	Penalty	The penalty object.
scheduler	HyperParamScheduler	The hyperparameter scheduler.

Example

from flyvision.utils.config_utils import get_default_config
# Note: the config is typically defined through the command line.
config = get_default_config(overrides=["task_name=flow",
                                       "ensemble_and_network_id=0"])
solver = MultiTaskSolver("test", config)
solver.train()

Source code in flyvision/solver.py

class MultiTaskSolver:
    """Implements training, testing, checkpoint, recovering of flyvis networks.

    Gives access to the network, decoder, task, optimizer, penalty and scheduler and
    the directory where the results are stored.

    Args:
        name: Name of the solver.
        config: Configuration for the solver.
        init_network: Whether to initialize the network. Defaults to True.
        init_decoder: Whether to initialize the decoder. Defaults to True.
        init_task: Whether to initialize the task. Defaults to True.
        init_optim: Whether to initialize the optimizer. Defaults to True.
        init_penalties: Whether to initialize penalties. Defaults to True.
        init_scheduler: Whether to initialize the scheduler. Defaults to True.
        delete_if_exists: Whether to delete existing directory. Defaults to False.

    Attributes:
        dir (NetworkDir): Directory where results are stored.
        network (Network): The neural network.
        decoder (Dict[str, nn.Module]): The decoder modules.
        task (Task): The task being solved.
        optimizer (torch.optim.Optimizer): The optimizer.
        penalty (Penalty): The penalty object.
        scheduler (HyperParamScheduler): The hyperparameter scheduler.

    Example:
        ```python
        from flyvision.utils.config_utils import get_default_config
        # Note: the config is typically defined through the command line.
        config = get_default_config(overrides=["task_name=flow",
                                               "ensemble_and_network_id=0"])
        solver = MultiTaskSolver("test", config)
        solver.train()
        ```
    """

    def __init__(
        self,
        name: str = "",
        config: Optional[Union[dict, Namespace]] = None,
        init_network: bool = True,
        init_decoder: bool = True,
        init_task: bool = True,
        init_optim: bool = True,
        init_penalties: bool = True,
        init_scheduler: bool = True,
        delete_if_exists: bool = False,
    ) -> None:
        name = name or config["network_name"]
        assert isinstance(name, str), "Provided name argument is not a string."
        self.dir = NetworkDir(
            name, {**(config or {}), **dict(delete_if_exists=delete_if_exists)}
        )

        self.path = self.dir.path

        self.config = self.dir.config

        self.iteration = 0
        self._val_loss = float("inf")
        self.checkpoint_path = self.dir.path / "chkpts"
        checkpoints = resolve_checkpoints(self.dir)
        self.checkpoints = checkpoints.indices
        self._last_chkpt_ind = -1
        self._curr_chkpt_ind = -1

        self._initialized = self._init_solver(
            init_network=init_network,
            init_decoder=init_decoder,
            init_task=init_task,
            init_optim=init_optim,
            init_penalties=init_penalties,
            init_scheduler=init_scheduler,
        )

        logging.info("Initialized solver.")
        logging.info(repr(self.config))

    def _init_solver(
        self,
        init_network: bool = False,
        init_decoder: bool = False,
        init_task: bool = False,
        init_optim: bool = False,
        init_penalties: bool = False,
        init_scheduler: bool = False,
    ) -> list:
        """Initialize solver components.

        Args:
            init_network: Whether to initialize the network.
            init_decoder: Whether to initialize the decoder.
            init_task: Whether to initialize the task.
            init_optim: Whether to initialize the optimizer.
            init_penalties: Whether to initialize penalties.
            init_scheduler: Whether to initialize the scheduler.

        Returns:
            A list of initialized components.
        """
        initialized = []

        if init_network:
            self.network = Network(**self.config.network)
            initialized.append("network")

        if init_task:
            self.task = Task(**self.config.task)
            initialized.append("task")

            if init_decoder:
                self.decoder = self.task.init_decoder(self.network.connectome)
                initialized.append("decoder")

        if init_optim:
            self.optimizer = self._init_optimizer(
                self.config.optim, self.network, self.decoder
            )
            initialized.append("optim")

        if init_penalties:
            self.penalty = Penalty(self.config.penalizer, self.network)
            initialized.append("penalties")

        if init_scheduler:
            self.scheduler = HyperParamScheduler(
                self.config.scheduler,
                self.network,
                self.task,
                self.optimizer,
                self.penalty,
            )
            self.scheduler(self.iteration)
            initialized.append("scheduler")

        return initialized

    @staticmethod
    def _init_optimizer(
        optim: Namespace, network: Network, decoder: Optional[Dict[str, nn.Module]]
    ) -> torch.optim.Optimizer:
        """Initializes the optimizer for network and decoder.

        Args:
            optim: Optimizer configuration.
            network: The neural network.
            decoder: The decoder modules.

        Returns:
            The initialized optimizer.
        """

        def decoder_parameters(decoder: Dict[str, nn.Module]):
            """Returns decoder parameters."""
            params = []
            for nn_module in decoder.values():
                params.append(
                    dict(
                        params=[w for w in nn_module.parameters()],
                        **config.optim_dec,
                    )
                )
            return params

        config = optim.deepcopy()

        optim_type = config.pop("type", "Adam")
        optim = torch.optim.__dict__[optim_type]
        logging.info("Initializing %s for network and decoder.", optim.__name__)

        param_groups = [dict(params=network.parameters(), **config.optim_net)]

        if decoder:
            param_groups.extend(decoder_parameters(decoder))

        return optim(param_groups)

    def train(self, overfit: bool = False, initial_checkpoint: bool = True) -> None:
        """Trains the network by backprop through time.

        Args:
            overfit: If true, the dataloader is substituted by a
                single-sequence loader and augmentation is turned off.
            initial_checkpoint: Whether to create an initial checkpoint when debugging.

        Raises:
            OverflowError: If the activity or loss reports NaN values for more
                than 100 iterations.

        Stores:
            ```bash
            dir / loss.h5
            dir / loss_<task>.h5
            dir / activity.h5
            dir / activity_min.h5
            dir / activity_max.h5
            ```
        """
        # return if iterations have already been trained.
        if self.iteration >= self.task.n_iters:
            return

        # to debug code within the training loop the initial checkpoint should be
        # disabled
        if initial_checkpoint:
            self.checkpoint()

        logging.info("Starting training.")
        # The overfit_data dataloader only contains a single sequence and
        # this is to debug the model architecture, configs etc.
        dataloader = self.task.overfit_data if overfit else self.task.train_data
        # For overfitting we also turn the augmentation off.
        augment = not overfit

        # The number of full presentations of the training data is derived from the
        # preset number of training iterations, the length of the dataloader and the
        # current iteration.
        n_epochs = np.ceil((self.task.n_iters - self.iteration) / len(dataloader)).astype(
            int
        )

        # This is after how many epochs the training states are checkpointed.
        chkpt_every_epoch = self.config.scheduler.chkpt_every_epoch

        logging.info("Training for %s epochs.", n_epochs)
        logging.info("Checkpointing every %s epochs.", chkpt_every_epoch)

        # Initialize data structures to store the loss and activity over iterations.
        loss_over_iters = []
        activity_over_iters = []
        activity_min_over_iters = []
        activity_max_over_iters = []
        loss_per_task = {f"loss_{task}": [] for task in self.task.dataset.tasks}

        start_time = time.time()
        with self.task.dataset.augmentation(augment):
            for epoch in range(n_epochs):
                # The default is to compute a steady state for each epoch, then
                # it's computed here. Note: unless done per iteration, parameter updates
                # within epochs are not considered in the steady state.
                steady_state = self.network.steady_state(
                    t_pre=self.config.get("t_pre_train", 0.5),
                    dt=self.task.dataset.dt,
                    batch_size=dataloader.batch_size,
                    value=0.5,
                )

                for _, data in enumerate(dataloader):

                    def handle_batch(data, steady_state):
                        """Closure to free memory by garbage collector effectively."""

                        # Resets the stimulus buffer (samples, frames, neurons).
                        n_samples, n_frames, _, _ = data["lum"].shape
                        self.network.stimulus.zero(n_samples, n_frames)

                        # Add batch of hex-videos (#frames, #samples, #hexals) as
                        # photorecptor stimuli.
                        self.network.stimulus.add_input(data["lum"])

                        # Reset gradients.
                        self.optimizer.zero_grad()

                        # Run stimulus through network.
                        activity = self.network(
                            self.network.stimulus(),
                            self.task.dataset.dt,
                            state=steady_state,
                        )

                        losses = {task: 0 for task in self.task.dataset.tasks}
                        for task in self.task.dataset.tasks:
                            y = data[task]
                            y_est = self.decoder[task](activity)

                            # to pass additional kwargs to the loss function, from
                            # the data batch from the dataset
                            losses[task] = self.task.loss(
                                y_est, y, task, **data.get("loss_kwargs", {})
                            )

                        # Sum all task losses. The weighting of the tasks is done in the
                        # loss function.
                        loss = sum(losses.values())

                        # Compute gradients.
                        loss.backward(retain_graph=True)
                        # Update parameters.
                        self.optimizer.step()

                        # Activity and parameter dependent penalties.
                        self.penalty(activity=activity, iteration=self.iteration)

                        # Log results.
                        loss = loss.detach().cpu()
                        for task in self.task.dataset.tasks:
                            loss_per_task[f"loss_{task}"].append(
                                losses[task].detach().cpu()
                            )
                        loss_over_iters.append(loss)
                        activity = activity.detach().cpu()
                        mean_activity = activity.mean()
                        activity_over_iters.append(mean_activity)
                        activity_min_over_iters.append(activity.min())
                        activity_max_over_iters.append(activity.max())
                        return loss, mean_activity

                    # Call closure.
                    loss, mean_activity = handle_batch(data, steady_state)

                    # Increment iteration count.
                    self.iteration += 1

                # Interrupt training if the network explodes.
                if torch.isnan(loss) or torch.isnan(mean_activity):
                    logging.warning("Network exploded.")
                    raise OverflowError("Invalid values encountered in trace.")

                # The scheduling of hyperparams are functions of the iteration
                # however, we allow steps only after full presentations of the data.
                if epoch + 1 != n_epochs:
                    self.scheduler(self.iteration)
                    logging.info("Scheduled paremeters for iteration %s.", self.iteration)

                # Checkpointing.
                if (epoch % chkpt_every_epoch == 0) or (epoch + 1 == n_epochs):
                    self.dir.loss = loss_over_iters
                    self.dir.activity = activity_over_iters
                    self.dir.activity_min = activity_min_over_iters
                    self.dir.activity_max = activity_max_over_iters

                    for task in self.task.dataset.tasks:
                        self.dir[f"loss_{task}"] = loss_per_task[f"loss_{task}"]

                    self.checkpoint()

                logging.info("Finished epoch.")

        time_elapsed = time.time() - start_time
        time_trained = self.dir.time_trained[()] if "time_trained" in self.dir else 0
        self.dir.time_trained = time_elapsed + time_trained
        logging.info("Finished training.")

    def checkpoint(self) -> None:
        """Creates a checkpoint.

        Validates on the validation data calling ~self.test.
        Validates on a training batch calling ~self.track_batch.
        Stores a checkpoint of the network, decoder and optimizer parameters using
        pytorch's pickle function.

        Stores:
            ```bash
            dir / chkpt_index.h5  # (List): numerical identifier of the checkpoint.
            dir / chkpt_iter.h5  # (List): iteration at which this checkpoint was
                                 # recorded.
            dir / best_chkpt_index.h5  # (int): chkpt index at which the val loss is
                                       # minimal.
            dir / dt.h5  # (float): the current time constant of the dataset.
            dir / chkpts / chkpt_<chkpt_index>  # (dict): the state dicts of the network,
                                                # decoder and optimizer.
            ```
        """
        self._last_chkpt_ind += 1
        self._curr_chkpt_ind += 1

        # Tracking of validation loss and training batch loss.
        logging.info("Test on validation data.")
        val_loss = self.test(
            dataloader=self.task.val_data, subdir="validation", track_loss=True
        )
        logging.info("Test on validation batch.")
        _ = self.test(
            dataloader=self.task.val_batch, subdir="validation_batch", track_loss=True
        )
        logging.info("Test on training data.")
        _ = self.test(dataloader=self.task.train_data, subdir="training", track_loss=True)
        logging.info("Test on training batch.")
        _ = self.test(
            dataloader=self.task.train_batch, subdir="training_batch", track_loss=True
        )

        logging.info("Saving state dicts.")
        # Store state of pytorch modules.
        nn_state_dict = self.network.state_dict()
        dec_state_dict = {}
        if self.decoder:
            dec_state_dict = valmap(lambda x: x.state_dict(), self.decoder)
        chkpt = {
            "network": nn_state_dict,
            "decoder": dec_state_dict,
            "optim": self.optimizer.state_dict(),
            "time": time.ctime(),
            "val_loss": val_loss,
            "iteration": self.iteration - 1,
            "dt": self.task.dataset.dt,
        }
        if hasattr(self, "penalty"):
            chkpt.update(self.penalty._chkpt())
        torch.save(chkpt, self.checkpoint_path / f"chkpt_{self._last_chkpt_ind:05}")

        # Append chkpt index.
        self.checkpoints.append(self._last_chkpt_ind)
        self.dir.extend("chkpt_index", [self._last_chkpt_ind])
        self.dir.extend("chkpt_iter", [self.iteration - 1])
        self.dir.dt = self.task.dataset.dt

        # Overwrite best val loss.
        if val_loss < self._val_loss:
            self.dir.best_chkpt_index = self._last_chkpt_ind
            self._val_loss = val_loss

        logging.info("Checkpointed.")

    @torch.no_grad()
    def test(
        self,
        dataloader: torch.utils.data.DataLoader,
        subdir: str = "validation",
        track_loss: bool = False,
        t_pre: float = 0.25,
    ) -> float:
        """Tests the network on a given dataloader.

        Args:
            dataloader: Data to test on.
            subdir: Name of subdirectory. Defaults to 'validation'.
            track_loss: Whether to store the loss in dir.subdir.
            t_pre: Warmup time before the stimulus starts.

        Returns:
            Validation loss.

        Stores:
            ```bash
            dir.<subdir>.loss_<task>  # (List): Loss per task, averaged over whole
                                      # dataset.
            dir.<subdir>.iteration  # (List): Iteration when this was called.
            dir.<subdir>.loss  # (List): Average loss over tasks.
            ```
        """
        self._eval()
        logging.info("Test")

        # Update hypterparams.
        self.scheduler(self.iteration)

        initial_state = self.network.steady_state(
            t_pre=t_pre,
            dt=self.task.dataset.dt,
            batch_size=dataloader.batch_size,
            value=0.5,
        )
        losses = {task: () for task in self.task.dataset.tasks}  # type: Dict[str, Tuple]

        with self.task.dataset.augmentation(False):
            for _, data in enumerate(dataloader):
                n_samples, n_frames, _, _ = data["lum"].shape
                self.network.stimulus.zero(n_samples, n_frames)

                self.network.stimulus.add_input(data["lum"])

                activity = self.network(
                    self.network.stimulus(),
                    self.task.dataset.dt,
                    state=initial_state,
                )

                for task in self.task.dataset.tasks:
                    y = data[task]
                    y_est = self.decoder[task](activity)

                    losses[task] += (
                        self.task.loss(y_est, y, task, **data.get("loss_kwargs", {}))
                        .detach()
                        .cpu()
                        .item(),
                    )

        # track loss per task.
        avg_loss_per_task = {}
        for task in self.task.dataset.tasks:
            # average the loss over the whole dataset
            avg_loss_per_task[task] = np.mean(losses[task])
            if track_loss:
                self.dir[subdir].extend("loss" + "_" + task, [avg_loss_per_task[task]])

        # average the loss over all tasks with equal weight
        summed_loss = sum(avg_loss_per_task.values())
        val_loss = summed_loss / len(avg_loss_per_task)

        if track_loss:
            self.dir[subdir].extend("iteration", [self.iteration])
            self.dir[subdir].extend("loss", [val_loss])

        self._train()

        return val_loss

    def _train(self) -> None:
        """Sets nn.Modules to train state."""
        self.network.train()
        if self.decoder is not None:
            for decoder in self.decoder.values():
                decoder.train()

    def _eval(self) -> None:
        """Sets nn.Modules to eval state."""
        self.network.eval()
        if self.decoder is not None:
            for decoder in self.decoder.values():
                decoder.eval()

    def recover(
        self,
        network: bool = True,
        decoder: bool = True,
        optimizer: bool = True,
        penalty: bool = True,
        checkpoint: Union[int, str] = "best",
        validation_subdir: str = "validation",
        loss_file_name: str = "loss",
        strict: bool = True,
        force: bool = False,
    ) -> None:
        """Recovers the solver state from a checkpoint.

        Args:
            network: Recover network parameters.
            decoder: Recover decoder parameters.
            optimizer: Recover optimizer parameters.
            penalty: Recover penalty parameters.
            checkpoint: Index of the checkpoint to recover.
            validation_subdir: Name of the subdir to base the best checkpoint on.
            loss_file_name: Name of the loss to base the best checkpoint on.
            strict: Whether to load the state dict of the decoders strictly.
            force: Force recovery of checkpoint if _curr_chkpt_ind is already
                the same as the checkpoint index.
        """
        checkpoints = resolve_checkpoints(
            self.dir, checkpoint, validation_subdir, loss_file_name
        )

        if checkpoint.index is None or not any((network, decoder, optimizer, penalty)):
            logging.info("No checkpoint found. Continuing with initialized parameters.")
            return

        if checkpoints.index == self._curr_chkpt_ind and not force:
            logging.info("Checkpoint already recovered.")
            return

        # Set the current and last checkpoint index. New checkpoints incrementally
        # increase the last checkpoint index.
        self._last_chkpt_ind = checkpoints.indices[-1]
        self._curr_chkpt_ind = checkpoints.index

        # Load checkpoint data.
        state_dict = torch.load(checkpoints.path)
        logging.info(f"Checkpoint {checkpoints.path} loaded.")

        self.iteration = state_dict.get("iteration", None)

        if "scheduler" in self._initialized:
            # Set the scheduler to the right iteration.
            self.scheduler(self.iteration)

        # The _val_loss variable is used to keep track of the best checkpoint according
        # to the evaluation routine during training.
        self._val_loss = state_dict.pop("val_loss", float("inf"))

        if network and "network" in self._initialized:
            recover_network(self.network, state_dict)
        if decoder and "decoder" in self._initialized:
            recover_decoder(self.decoder, state_dict, strict=strict)
        if optimizer and "optim" in self._initialized:
            recover_optimizer(self.optimizer, state_dict)
        if penalty and "penalties" in self._initialized:
            recover_penalty_optimizers(self.penalty.optimizers, state_dict)

        logging.info("Recovered modules.")

train

train(overfit=False, initial_checkpoint=True)

Trains the network by backprop through time.

Parameters:

Name	Type	Description	Default
overfit	bool	If true, the dataloader is substituted by a single-sequence loader and augmentation is turned off.	False
initial_checkpoint	bool	Whether to create an initial checkpoint when debugging.	True

Raises:

Type	Description
OverflowError	If the activity or loss reports NaN values for more than 100 iterations.

Stores

dir / loss.h5
dir / loss_<task>.h5
dir / activity.h5
dir / activity_min.h5
dir / activity_max.h5

Source code in flyvision/solver.py

def train(self, overfit: bool = False, initial_checkpoint: bool = True) -> None:
    """Trains the network by backprop through time.

    Args:
        overfit: If true, the dataloader is substituted by a
            single-sequence loader and augmentation is turned off.
        initial_checkpoint: Whether to create an initial checkpoint when debugging.

    Raises:
        OverflowError: If the activity or loss reports NaN values for more
            than 100 iterations.

    Stores:
        ```bash
        dir / loss.h5
        dir / loss_<task>.h5
        dir / activity.h5
        dir / activity_min.h5
        dir / activity_max.h5
        ```
    """
    # return if iterations have already been trained.
    if self.iteration >= self.task.n_iters:
        return

    # to debug code within the training loop the initial checkpoint should be
    # disabled
    if initial_checkpoint:
        self.checkpoint()

    logging.info("Starting training.")
    # The overfit_data dataloader only contains a single sequence and
    # this is to debug the model architecture, configs etc.
    dataloader = self.task.overfit_data if overfit else self.task.train_data
    # For overfitting we also turn the augmentation off.
    augment = not overfit

    # The number of full presentations of the training data is derived from the
    # preset number of training iterations, the length of the dataloader and the
    # current iteration.
    n_epochs = np.ceil((self.task.n_iters - self.iteration) / len(dataloader)).astype(
        int
    )

    # This is after how many epochs the training states are checkpointed.
    chkpt_every_epoch = self.config.scheduler.chkpt_every_epoch

    logging.info("Training for %s epochs.", n_epochs)
    logging.info("Checkpointing every %s epochs.", chkpt_every_epoch)

    # Initialize data structures to store the loss and activity over iterations.
    loss_over_iters = []
    activity_over_iters = []
    activity_min_over_iters = []
    activity_max_over_iters = []
    loss_per_task = {f"loss_{task}": [] for task in self.task.dataset.tasks}

    start_time = time.time()
    with self.task.dataset.augmentation(augment):
        for epoch in range(n_epochs):
            # The default is to compute a steady state for each epoch, then
            # it's computed here. Note: unless done per iteration, parameter updates
            # within epochs are not considered in the steady state.
            steady_state = self.network.steady_state(
                t_pre=self.config.get("t_pre_train", 0.5),
                dt=self.task.dataset.dt,
                batch_size=dataloader.batch_size,
                value=0.5,
            )

            for _, data in enumerate(dataloader):

                def handle_batch(data, steady_state):
                    """Closure to free memory by garbage collector effectively."""

                    # Resets the stimulus buffer (samples, frames, neurons).
                    n_samples, n_frames, _, _ = data["lum"].shape
                    self.network.stimulus.zero(n_samples, n_frames)

                    # Add batch of hex-videos (#frames, #samples, #hexals) as
                    # photorecptor stimuli.
                    self.network.stimulus.add_input(data["lum"])

                    # Reset gradients.
                    self.optimizer.zero_grad()

                    # Run stimulus through network.
                    activity = self.network(
                        self.network.stimulus(),
                        self.task.dataset.dt,
                        state=steady_state,
                    )

                    losses = {task: 0 for task in self.task.dataset.tasks}
                    for task in self.task.dataset.tasks:
                        y = data[task]
                        y_est = self.decoder[task](activity)

                        # to pass additional kwargs to the loss function, from
                        # the data batch from the dataset
                        losses[task] = self.task.loss(
                            y_est, y, task, **data.get("loss_kwargs", {})
                        )

                    # Sum all task losses. The weighting of the tasks is done in the
                    # loss function.
                    loss = sum(losses.values())

                    # Compute gradients.
                    loss.backward(retain_graph=True)
                    # Update parameters.
                    self.optimizer.step()

                    # Activity and parameter dependent penalties.
                    self.penalty(activity=activity, iteration=self.iteration)

                    # Log results.
                    loss = loss.detach().cpu()
                    for task in self.task.dataset.tasks:
                        loss_per_task[f"loss_{task}"].append(
                            losses[task].detach().cpu()
                        )
                    loss_over_iters.append(loss)
                    activity = activity.detach().cpu()
                    mean_activity = activity.mean()
                    activity_over_iters.append(mean_activity)
                    activity_min_over_iters.append(activity.min())
                    activity_max_over_iters.append(activity.max())
                    return loss, mean_activity

                # Call closure.
                loss, mean_activity = handle_batch(data, steady_state)

                # Increment iteration count.
                self.iteration += 1

            # Interrupt training if the network explodes.
            if torch.isnan(loss) or torch.isnan(mean_activity):
                logging.warning("Network exploded.")
                raise OverflowError("Invalid values encountered in trace.")

            # The scheduling of hyperparams are functions of the iteration
            # however, we allow steps only after full presentations of the data.
            if epoch + 1 != n_epochs:
                self.scheduler(self.iteration)
                logging.info("Scheduled paremeters for iteration %s.", self.iteration)

            # Checkpointing.
            if (epoch % chkpt_every_epoch == 0) or (epoch + 1 == n_epochs):
                self.dir.loss = loss_over_iters
                self.dir.activity = activity_over_iters
                self.dir.activity_min = activity_min_over_iters
                self.dir.activity_max = activity_max_over_iters

                for task in self.task.dataset.tasks:
                    self.dir[f"loss_{task}"] = loss_per_task[f"loss_{task}"]

                self.checkpoint()

            logging.info("Finished epoch.")

    time_elapsed = time.time() - start_time
    time_trained = self.dir.time_trained[()] if "time_trained" in self.dir else 0
    self.dir.time_trained = time_elapsed + time_trained
    logging.info("Finished training.")

checkpoint

checkpoint()

Creates a checkpoint.

Validates on the validation data calling ~self.test. Validates on a training batch calling ~self.track_batch. Stores a checkpoint of the network, decoder and optimizer parameters using pytorch’s pickle function.

Stores

dir / chkpt_index.h5  # (List): numerical identifier of the checkpoint.
dir / chkpt_iter.h5  # (List): iteration at which this checkpoint was
                     # recorded.
dir / best_chkpt_index.h5  # (int): chkpt index at which the val loss is
                           # minimal.
dir / dt.h5  # (float): the current time constant of the dataset.
dir / chkpts / chkpt_<chkpt_index>  # (dict): the state dicts of the network,
                                    # decoder and optimizer.

Source code in flyvision/solver.py

def checkpoint(self) -> None:
    """Creates a checkpoint.

    Validates on the validation data calling ~self.test.
    Validates on a training batch calling ~self.track_batch.
    Stores a checkpoint of the network, decoder and optimizer parameters using
    pytorch's pickle function.

    Stores:
        ```bash
        dir / chkpt_index.h5  # (List): numerical identifier of the checkpoint.
        dir / chkpt_iter.h5  # (List): iteration at which this checkpoint was
                             # recorded.
        dir / best_chkpt_index.h5  # (int): chkpt index at which the val loss is
                                   # minimal.
        dir / dt.h5  # (float): the current time constant of the dataset.
        dir / chkpts / chkpt_<chkpt_index>  # (dict): the state dicts of the network,
                                            # decoder and optimizer.
        ```
    """
    self._last_chkpt_ind += 1
    self._curr_chkpt_ind += 1

    # Tracking of validation loss and training batch loss.
    logging.info("Test on validation data.")
    val_loss = self.test(
        dataloader=self.task.val_data, subdir="validation", track_loss=True
    )
    logging.info("Test on validation batch.")
    _ = self.test(
        dataloader=self.task.val_batch, subdir="validation_batch", track_loss=True
    )
    logging.info("Test on training data.")
    _ = self.test(dataloader=self.task.train_data, subdir="training", track_loss=True)
    logging.info("Test on training batch.")
    _ = self.test(
        dataloader=self.task.train_batch, subdir="training_batch", track_loss=True
    )

    logging.info("Saving state dicts.")
    # Store state of pytorch modules.
    nn_state_dict = self.network.state_dict()
    dec_state_dict = {}
    if self.decoder:
        dec_state_dict = valmap(lambda x: x.state_dict(), self.decoder)
    chkpt = {
        "network": nn_state_dict,
        "decoder": dec_state_dict,
        "optim": self.optimizer.state_dict(),
        "time": time.ctime(),
        "val_loss": val_loss,
        "iteration": self.iteration - 1,
        "dt": self.task.dataset.dt,
    }
    if hasattr(self, "penalty"):
        chkpt.update(self.penalty._chkpt())
    torch.save(chkpt, self.checkpoint_path / f"chkpt_{self._last_chkpt_ind:05}")

    # Append chkpt index.
    self.checkpoints.append(self._last_chkpt_ind)
    self.dir.extend("chkpt_index", [self._last_chkpt_ind])
    self.dir.extend("chkpt_iter", [self.iteration - 1])
    self.dir.dt = self.task.dataset.dt

    # Overwrite best val loss.
    if val_loss < self._val_loss:
        self.dir.best_chkpt_index = self._last_chkpt_ind
        self._val_loss = val_loss

    logging.info("Checkpointed.")

test

test(
    dataloader,
    subdir="validation",
    track_loss=False,
    t_pre=0.25,
)

Tests the network on a given dataloader.

Parameters:

Name	Type	Description	Default
dataloader	DataLoader	Data to test on.	required
subdir	str	Name of subdirectory. Defaults to ‘validation’.	'validation'
track_loss	bool	Whether to store the loss in dir.subdir.	False
t_pre	float	Warmup time before the stimulus starts.	0.25

Returns:

Type	Description
float	Validation loss.

Stores

dir.<subdir>.loss_<task>  # (List): Loss per task, averaged over whole
                          # dataset.
dir.<subdir>.iteration  # (List): Iteration when this was called.
dir.<subdir>.loss  # (List): Average loss over tasks.

Source code in flyvision/solver.py

@torch.no_grad()
def test(
    self,
    dataloader: torch.utils.data.DataLoader,
    subdir: str = "validation",
    track_loss: bool = False,
    t_pre: float = 0.25,
) -> float:
    """Tests the network on a given dataloader.

    Args:
        dataloader: Data to test on.
        subdir: Name of subdirectory. Defaults to 'validation'.
        track_loss: Whether to store the loss in dir.subdir.
        t_pre: Warmup time before the stimulus starts.

    Returns:
        Validation loss.

    Stores:
        ```bash
        dir.<subdir>.loss_<task>  # (List): Loss per task, averaged over whole
                                  # dataset.
        dir.<subdir>.iteration  # (List): Iteration when this was called.
        dir.<subdir>.loss  # (List): Average loss over tasks.
        ```
    """
    self._eval()
    logging.info("Test")

    # Update hypterparams.
    self.scheduler(self.iteration)

    initial_state = self.network.steady_state(
        t_pre=t_pre,
        dt=self.task.dataset.dt,
        batch_size=dataloader.batch_size,
        value=0.5,
    )
    losses = {task: () for task in self.task.dataset.tasks}  # type: Dict[str, Tuple]

    with self.task.dataset.augmentation(False):
        for _, data in enumerate(dataloader):
            n_samples, n_frames, _, _ = data["lum"].shape
            self.network.stimulus.zero(n_samples, n_frames)

            self.network.stimulus.add_input(data["lum"])

            activity = self.network(
                self.network.stimulus(),
                self.task.dataset.dt,
                state=initial_state,
            )

            for task in self.task.dataset.tasks:
                y = data[task]
                y_est = self.decoder[task](activity)

                losses[task] += (
                    self.task.loss(y_est, y, task, **data.get("loss_kwargs", {}))
                    .detach()
                    .cpu()
                    .item(),
                )

    # track loss per task.
    avg_loss_per_task = {}
    for task in self.task.dataset.tasks:
        # average the loss over the whole dataset
        avg_loss_per_task[task] = np.mean(losses[task])
        if track_loss:
            self.dir[subdir].extend("loss" + "_" + task, [avg_loss_per_task[task]])

    # average the loss over all tasks with equal weight
    summed_loss = sum(avg_loss_per_task.values())
    val_loss = summed_loss / len(avg_loss_per_task)

    if track_loss:
        self.dir[subdir].extend("iteration", [self.iteration])
        self.dir[subdir].extend("loss", [val_loss])

    self._train()

    return val_loss

recover

recover(
    network=True,
    decoder=True,
    optimizer=True,
    penalty=True,
    checkpoint="best",
    validation_subdir="validation",
    loss_file_name="loss",
    strict=True,
    force=False,
)

Recovers the solver state from a checkpoint.

Parameters:

Name	Type	Description	Default
network	bool	Recover network parameters.	True
decoder	bool	Recover decoder parameters.	True
optimizer	bool	Recover optimizer parameters.	True
penalty	bool	Recover penalty parameters.	True
checkpoint	Union[int, str]	Index of the checkpoint to recover.	'best'
validation_subdir	str	Name of the subdir to base the best checkpoint on.	'validation'
loss_file_name	str	Name of the loss to base the best checkpoint on.	'loss'
strict	bool	Whether to load the state dict of the decoders strictly.	True
force	bool	Force recovery of checkpoint if _curr_chkpt_ind is already the same as the checkpoint index.	False

Source code in flyvision/solver.py

def recover(
    self,
    network: bool = True,
    decoder: bool = True,
    optimizer: bool = True,
    penalty: bool = True,
    checkpoint: Union[int, str] = "best",
    validation_subdir: str = "validation",
    loss_file_name: str = "loss",
    strict: bool = True,
    force: bool = False,
) -> None:
    """Recovers the solver state from a checkpoint.

    Args:
        network: Recover network parameters.
        decoder: Recover decoder parameters.
        optimizer: Recover optimizer parameters.
        penalty: Recover penalty parameters.
        checkpoint: Index of the checkpoint to recover.
        validation_subdir: Name of the subdir to base the best checkpoint on.
        loss_file_name: Name of the loss to base the best checkpoint on.
        strict: Whether to load the state dict of the decoders strictly.
        force: Force recovery of checkpoint if _curr_chkpt_ind is already
            the same as the checkpoint index.
    """
    checkpoints = resolve_checkpoints(
        self.dir, checkpoint, validation_subdir, loss_file_name
    )

    if checkpoint.index is None or not any((network, decoder, optimizer, penalty)):
        logging.info("No checkpoint found. Continuing with initialized parameters.")
        return

    if checkpoints.index == self._curr_chkpt_ind and not force:
        logging.info("Checkpoint already recovered.")
        return

    # Set the current and last checkpoint index. New checkpoints incrementally
    # increase the last checkpoint index.
    self._last_chkpt_ind = checkpoints.indices[-1]
    self._curr_chkpt_ind = checkpoints.index

    # Load checkpoint data.
    state_dict = torch.load(checkpoints.path)
    logging.info(f"Checkpoint {checkpoints.path} loaded.")

    self.iteration = state_dict.get("iteration", None)

    if "scheduler" in self._initialized:
        # Set the scheduler to the right iteration.
        self.scheduler(self.iteration)

    # The _val_loss variable is used to keep track of the best checkpoint according
    # to the evaluation routine during training.
    self._val_loss = state_dict.pop("val_loss", float("inf"))

    if network and "network" in self._initialized:
        recover_network(self.network, state_dict)
    if decoder and "decoder" in self._initialized:
        recover_decoder(self.decoder, state_dict, strict=strict)
    if optimizer and "optim" in self._initialized:
        recover_optimizer(self.optimizer, state_dict)
    if penalty and "penalties" in self._initialized:
        recover_penalty_optimizers(self.penalty.optimizers, state_dict)

    logging.info("Recovered modules.")

Penalty

Penalties on specific parameters.

Parameters:

Name	Type	Description	Default
penalty	Namespace	Penalty configuration.	required
network	Network	The neural network.	required

Default config in config/penalizer/penalizer.yaml

activity_penalty:
    activity_baseline: 5.0
    activity_penalty: 0.1
    stop_iter: 150000
    below_baseline_penalty_weight: 1.0
    above_baseline_penalty_weight: 0.1
optim: SGD

Default config in config/network/node_config/bias/bias.yaml

type: RestingPotential
groupby:
    - type
initial_dist: Normal
mode: sample
requires_grad: true
seed: 0
mean: 0.5
std: 0.05
symmetric: []
penalize:
    activity: true

Attributes:

Name	Type	Description
config	Namespace	Penalty configuration.
network	Network	The neural network.
central_cells_index	ndarray	Index of central cells.
parameter_config	Namespace	Configuration for parameter penalties.
activity_penalty	float	Penalty for activity.
activity_baseline	float	Baseline for activity.
activity_penalty_stop_iter	int	Iteration to stop activity penalty.
below_baseline_penalty_weight	float	Weight for below baseline penalty.
above_baseline_penalty_weight	float	Weight for above baseline penalty.
parameter_optim	Optimizer	Optimizer for parameter penalties.
activity_optim	Optimizer	Optimizer for activity penalties.
optimizers	Dict[str, Optimizer]	Dictionary of optimizers.
param_list_func_pen	list	List of parameters for function penalties.
param_list_act_pen	list	List of parameters for activity penalties.

Examples:

Example configurations passed to the network object:

# Example 1: Penalize the resting potential of all cell types.
bias = Namespace(
    ... (other parameters)
    penalize=Namespace(activity=True),
)

# Example 2: add a weight decay penalty to all synapse strengths.
syn_strength = Namespace(
    ... (other parameters)
    penalize=Namespace(function="weight_decay", kwargs=dict(lambda=1e-3,)),
)

Source code in flyvision/solver.py

class Penalty:
    """Penalties on specific parameters.

    Args:
        penalty: Penalty configuration.
        network: The neural network.

    Note: Default config in config/penalizer/penalizer.yaml
        ```yaml
        activity_penalty:
            activity_baseline: 5.0
            activity_penalty: 0.1
            stop_iter: 150000
            below_baseline_penalty_weight: 1.0
            above_baseline_penalty_weight: 0.1
        optim: SGD
        ```

    Note: Default config in config/network/node_config/bias/bias.yaml
        ```yaml
        type: RestingPotential
        groupby:
            - type
        initial_dist: Normal
        mode: sample
        requires_grad: true
        seed: 0
        mean: 0.5
        std: 0.05
        symmetric: []
        penalize:
            activity: true
        ```

    Attributes:
        config (Namespace): Penalty configuration.
        network (Network): The neural network.
        central_cells_index (np.ndarray): Index of central cells.
        parameter_config (Namespace): Configuration for parameter penalties.
        activity_penalty (float): Penalty for activity.
        activity_baseline (float): Baseline for activity.
        activity_penalty_stop_iter (int): Iteration to stop activity penalty.
        below_baseline_penalty_weight (float): Weight for below baseline penalty.
        above_baseline_penalty_weight (float): Weight for above baseline penalty.
        parameter_optim (torch.optim.Optimizer): Optimizer for parameter penalties.
        activity_optim (torch.optim.Optimizer): Optimizer for activity penalties.
        optimizers (Dict[str, torch.optim.Optimizer]): Dictionary of optimizers.
        param_list_func_pen (list): List of parameters for function penalties.
        param_list_act_pen (list): List of parameters for activity penalties.

    Examples:
        Example configurations passed to the network object:

        ```python
        # Example 1: Penalize the resting potential of all cell types.
        bias = Namespace(
            ... (other parameters)
            penalize=Namespace(activity=True),
        )

        # Example 2: add a weight decay penalty to all synapse strengths.
        syn_strength = Namespace(
            ... (other parameters)
            penalize=Namespace(function="weight_decay", kwargs=dict(lambda=1e-3,)),
        )
        ```
    """

    def __init__(self, penalty: Namespace, network: Network):
        self.config = penalty
        self.network = network
        self.central_cells_index = self.network.connectome.central_cells_index[:]

        self.parameter_config = self.get_network_param_penalty_configs()
        self.optim_class = getattr(torch.optim, getattr(self.config, "optim", "SGD"))
        self.parameter_optim, self.activity_optim = None, None
        self.init_optim()
        self.init_hparams()

    def get_network_param_penalty_configs(self) -> Namespace:
        """Returns a dictionary of all network parameters configured to be penalized."""
        node_config = Namespace({
            "nodes_" + k: v.pop("penalize", None)
            for k, v in self.network.config.node_config.deepcopy().items()
        })
        edge_config = Namespace({
            "edges_" + k: v.pop("penalize", None)
            for k, v in self.network.config.edge_config.deepcopy().items()
        })
        return valfilter(
            lambda v: v is not None,
            Namespace(**node_config, **edge_config),
            factory=Namespace,
        )

    def init_optim(self) -> None:
        """Initialize the individual optimizer instances with the correct set of
        parameters.
        """
        self.optimizers = {}
        self.param_list_func_pen = []
        self.param_list_act_pen = []

        # collect the parameters that need to be penalized
        # either by a function or by activity
        for name, config in self.parameter_config.items():
            if "function" in config and any(list(config.kwargs.values())):
                self.param_list_func_pen.append(name)
            if getattr(config, "activity", False):
                self.param_list_act_pen.append(name)

        if self.param_list_func_pen:
            self.parameter_optim = self.optim_class(
                (getattr(self.network, param) for param in self.param_list_func_pen),
                lr=1e-3,
            )  # LR is overwritten by scheduler.
            self.optimizers.update(dict(parameter_optim=self.parameter_optim))

        if self.param_list_act_pen:
            self.activity_optim = self.optim_class(
                (getattr(self.network, param) for param in self.param_list_act_pen),
                lr=1e-3,
            )  # LR is overwritten by scheduler.
            self.optimizers.update(dict(activity_optim=self.activity_optim))

    def init_hparams(self) -> None:
        """Initialize the hyperparameters for the activity penalty."""
        config = self.config.get("activity_penalty", Namespace())

        # collecting activity penalty parameters
        (
            self.activity_penalty,
            self.activity_baseline,
            self.activity_penalty_stop_iter,
            self.below_baseline_penalty_weight,
            self.above_baseline_penalty_weight,
        ) = (
            config.get("activity_penalty", None),
            config.get("activity_baseline", None),
            config.get("stop_iter", None),
            config.get("below_baseline_penalty_weight", None),
            config.get("above_baseline_penalty_weight", None),
        )

        if (
            not any((
                self.activity_penalty,
                self.activity_baseline,
                self.below_baseline_penalty_weight,
                self.above_baseline_penalty_weight,
            ))
            and self.param_list_act_pen
        ):
            raise ValueError(
                "Activity penalty is enabled but no activity penalty parameters are "
                "set."
            )

    def __repr__(self):
        return (
            f"Penalty("
            f"parameter_config={self.parameter_config}, "
            f"activity_penalty={self.activity_penalty}, "
            f"activity_baseline={self.activity_baseline}, "
            f"activity_penalty_stop_iter={self.activity_penalty_stop_iter}, "
            f"below_baseline_penalty_weight={self.below_baseline_penalty_weight}, "
            f"above_baseline_penalty_weight={self.above_baseline_penalty_weight}, "
            f"optim_class={self.optim_class}"
            f")"
        )

    def __call__(self, activity: torch.Tensor, iteration: int) -> None:
        """Run all configured penalties.

        Args:
            activity: Network activity.
            iteration: Current iteration.
        """
        if self.parameter_optim:
            self.param_penalty_step()
        if self.activity_optim:
            if (
                self.activity_penalty_stop_iter is None
                or iteration < self.activity_penalty_stop_iter
            ):
                self.activity_penalty_step(activity, retain_graph=False)
            else:
                self.activity_optim = None

    def _chkpt(self) -> dict:
        """Returns a dictionary of all state dicts of all optimizer instances."""
        _chkpt = {}
        for key, optim in self.optimizers.items():
            if optim is not None:
                _chkpt[key] = optim.state_dict()
        return _chkpt

    def param_penalty_step(self) -> None:
        """Apply all the penalties on the individual parameters."""
        self.parameter_optim.zero_grad()
        penalty = 0
        for param, config in self.parameter_config.items():
            if getattr(config, "function", False):
                penalty += getattr(self, config.function)(param, config)
        penalty.backward()
        self.parameter_optim.step()
        self.network.clamp()

    def activity_penalty_step(
        self, activity: torch.Tensor, retain_graph: bool = True
    ) -> None:
        """Penalizes parameters tracked in activity_optim for too high or low activity.

        Encourages the nodes to have a higher or lower temporal mean activity,
        remedying dead or overactive neurons.

        Note:
            This assumes that the central cells are representative for all cells because
            of shared parameters across the cell types, which makes this reasonably
            efficient.

        Args:
            activity: Network activity of shape (n_samples, n_frames, n_nodes).
            retain_graph: Whether to retain the computation graph.
        """
        self.activity_optim.zero_grad()
        n_samples, n_frames, n_nodes = activity.shape
        # the temporal average activity of the central nodes after a couple of frames
        # to avoid the initial transient response
        activity_mean = activity[:, n_frames // 4 :, self.central_cells_index].mean(
            dim=1
        )  # (n_samples, n_node_types)
        penalty = (
            self.activity_penalty
            * (
                asymmetric_weighting(
                    self.activity_baseline - activity_mean,
                    self.below_baseline_penalty_weight,
                    self.above_baseline_penalty_weight,
                )
                ** 2
            ).mean()
        )
        penalty.backward(retain_graph=retain_graph)
        self.activity_optim.step()
        self.network.clamp()

    def weight_decay(self, param: str, config: Namespace) -> torch.Tensor:
        """Adds weight decay to the loss.

        Warning: Experimental

        Args:
            param: Name of the parameter.
            config: Configuration for the penalty.

        Returns:
            The weight decay penalty.

        """
        w = getattr(self.network, param)
        return config.kwargs["lambda"] * (w**2).sum()

    def prior(self, param: str, config: Namespace) -> torch.Tensor:
        """L2 penalty towards initial values.

        Warning: Experimental

        Args:
            param: Name of the parameter.
            config: Configuration for the penalty.

        Returns:
            The L2 penalty.

        TODO: this might be a convenient but suboptimal implementation when the initial
        values are cast to tensors at each iteration.
        """
        _key = "edge_config" if param.startswith("edges") else "node_config"
        prior = torch.tensor(
            getattr(self.network.config, _key)[
                param.replace("edges_", "").replace("nodes_", "")
            ].value,
            dtype=torch.float32,
        )
        return (
            config.kwargs["lambda"] * ((getattr(self.network, param) - prior) ** 2).sum()
        )

get_network_param_penalty_configs

get_network_param_penalty_configs()

Returns a dictionary of all network parameters configured to be penalized.

Source code in flyvision/solver.py

def get_network_param_penalty_configs(self) -> Namespace:
    """Returns a dictionary of all network parameters configured to be penalized."""
    node_config = Namespace({
        "nodes_" + k: v.pop("penalize", None)
        for k, v in self.network.config.node_config.deepcopy().items()
    })
    edge_config = Namespace({
        "edges_" + k: v.pop("penalize", None)
        for k, v in self.network.config.edge_config.deepcopy().items()
    })
    return valfilter(
        lambda v: v is not None,
        Namespace(**node_config, **edge_config),
        factory=Namespace,
    )

init_optim

init_optim()

Initialize the individual optimizer instances with the correct set of parameters.

Source code in flyvision/solver.py

def init_optim(self) -> None:
    """Initialize the individual optimizer instances with the correct set of
    parameters.
    """
    self.optimizers = {}
    self.param_list_func_pen = []
    self.param_list_act_pen = []

    # collect the parameters that need to be penalized
    # either by a function or by activity
    for name, config in self.parameter_config.items():
        if "function" in config and any(list(config.kwargs.values())):
            self.param_list_func_pen.append(name)
        if getattr(config, "activity", False):
            self.param_list_act_pen.append(name)

    if self.param_list_func_pen:
        self.parameter_optim = self.optim_class(
            (getattr(self.network, param) for param in self.param_list_func_pen),
            lr=1e-3,
        )  # LR is overwritten by scheduler.
        self.optimizers.update(dict(parameter_optim=self.parameter_optim))

    if self.param_list_act_pen:
        self.activity_optim = self.optim_class(
            (getattr(self.network, param) for param in self.param_list_act_pen),
            lr=1e-3,
        )  # LR is overwritten by scheduler.
        self.optimizers.update(dict(activity_optim=self.activity_optim))

init_hparams

init_hparams()

Initialize the hyperparameters for the activity penalty.

Source code in flyvision/solver.py

def init_hparams(self) -> None:
    """Initialize the hyperparameters for the activity penalty."""
    config = self.config.get("activity_penalty", Namespace())

    # collecting activity penalty parameters
    (
        self.activity_penalty,
        self.activity_baseline,
        self.activity_penalty_stop_iter,
        self.below_baseline_penalty_weight,
        self.above_baseline_penalty_weight,
    ) = (
        config.get("activity_penalty", None),
        config.get("activity_baseline", None),
        config.get("stop_iter", None),
        config.get("below_baseline_penalty_weight", None),
        config.get("above_baseline_penalty_weight", None),
    )

    if (
        not any((
            self.activity_penalty,
            self.activity_baseline,
            self.below_baseline_penalty_weight,
            self.above_baseline_penalty_weight,
        ))
        and self.param_list_act_pen
    ):
        raise ValueError(
            "Activity penalty is enabled but no activity penalty parameters are "
            "set."
        )

__call__

__call__(activity, iteration)

Run all configured penalties.

Parameters:

Name	Type	Description	Default
activity	Tensor	Network activity.	required
iteration	int	Current iteration.	required

Source code in flyvision/solver.py

def __call__(self, activity: torch.Tensor, iteration: int) -> None:
    """Run all configured penalties.

    Args:
        activity: Network activity.
        iteration: Current iteration.
    """
    if self.parameter_optim:
        self.param_penalty_step()
    if self.activity_optim:
        if (
            self.activity_penalty_stop_iter is None
            or iteration < self.activity_penalty_stop_iter
        ):
            self.activity_penalty_step(activity, retain_graph=False)
        else:
            self.activity_optim = None

param_penalty_step

param_penalty_step()

Apply all the penalties on the individual parameters.

Source code in flyvision/solver.py

def param_penalty_step(self) -> None:
    """Apply all the penalties on the individual parameters."""
    self.parameter_optim.zero_grad()
    penalty = 0
    for param, config in self.parameter_config.items():
        if getattr(config, "function", False):
            penalty += getattr(self, config.function)(param, config)
    penalty.backward()
    self.parameter_optim.step()
    self.network.clamp()

activity_penalty_step

activity_penalty_step(activity, retain_graph=True)

Penalizes parameters tracked in activity_optim for too high or low activity.

Encourages the nodes to have a higher or lower temporal mean activity, remedying dead or overactive neurons.

Note

This assumes that the central cells are representative for all cells because of shared parameters across the cell types, which makes this reasonably efficient.

Parameters:

Name	Type	Description	Default
activity	Tensor	Network activity of shape (n_samples, n_frames, n_nodes).	required
retain_graph	bool	Whether to retain the computation graph.	True

Source code in flyvision/solver.py

def activity_penalty_step(
    self, activity: torch.Tensor, retain_graph: bool = True
) -> None:
    """Penalizes parameters tracked in activity_optim for too high or low activity.

    Encourages the nodes to have a higher or lower temporal mean activity,
    remedying dead or overactive neurons.

    Note:
        This assumes that the central cells are representative for all cells because
        of shared parameters across the cell types, which makes this reasonably
        efficient.

    Args:
        activity: Network activity of shape (n_samples, n_frames, n_nodes).
        retain_graph: Whether to retain the computation graph.
    """
    self.activity_optim.zero_grad()
    n_samples, n_frames, n_nodes = activity.shape
    # the temporal average activity of the central nodes after a couple of frames
    # to avoid the initial transient response
    activity_mean = activity[:, n_frames // 4 :, self.central_cells_index].mean(
        dim=1
    )  # (n_samples, n_node_types)
    penalty = (
        self.activity_penalty
        * (
            asymmetric_weighting(
                self.activity_baseline - activity_mean,
                self.below_baseline_penalty_weight,
                self.above_baseline_penalty_weight,
            )
            ** 2
        ).mean()
    )
    penalty.backward(retain_graph=retain_graph)
    self.activity_optim.step()
    self.network.clamp()

weight_decay

weight_decay(param, config)

Adds weight decay to the loss.

Warning: Experimental

Parameters:

Name	Type	Description	Default
param	str	Name of the parameter.	required
config	Namespace	Configuration for the penalty.	required

Returns:

Type	Description
Tensor	The weight decay penalty.

Source code in flyvision/solver.py

def weight_decay(self, param: str, config: Namespace) -> torch.Tensor:
    """Adds weight decay to the loss.

    Warning: Experimental

    Args:
        param: Name of the parameter.
        config: Configuration for the penalty.

    Returns:
        The weight decay penalty.

    """
    w = getattr(self.network, param)
    return config.kwargs["lambda"] * (w**2).sum()

prior

prior(param, config)

L2 penalty towards initial values.

Warning: Experimental

Parameters:

Name	Type	Description	Default
param	str	Name of the parameter.	required
config	Namespace	Configuration for the penalty.	required

Returns:

Type	Description
Tensor	The L2 penalty.

TODO: this might be a convenient but suboptimal implementation when the initial values are cast to tensors at each iteration.

Source code in flyvision/solver.py

def prior(self, param: str, config: Namespace) -> torch.Tensor:
    """L2 penalty towards initial values.

    Warning: Experimental

    Args:
        param: Name of the parameter.
        config: Configuration for the penalty.

    Returns:
        The L2 penalty.

    TODO: this might be a convenient but suboptimal implementation when the initial
    values are cast to tensors at each iteration.
    """
    _key = "edge_config" if param.startswith("edges") else "node_config"
    prior = torch.tensor(
        getattr(self.network.config, _key)[
            param.replace("edges_", "").replace("nodes_", "")
        ].value,
        dtype=torch.float32,
    )
    return (
        config.kwargs["lambda"] * ((getattr(self.network, param) - prior) ** 2).sum()
    )

HyperParamScheduler

Schedules hyperparameters per training iteration.

Calling the scheduler instance updates the respective hyperparameters per training iteration.

Parameters:

Name	Type	Description	Default
scheduler	Namespace	Scheduler configuration.	required
network	Optional[Network]	The neural network.	required
task	Optional[Task]	The task being solved.	required
optimizer	Optional[Optimizer]	The optimizer.	required
penalizer	Optional[Penalty]	The penalty object.	required

Default config in config/scheduler/scheduler.yaml

lr_net:
    function: stepwise
    start: 5.0e-05
    stop: 5.0e-06
    steps: 10
lr_dec:
    function: stepwise
    start: 5.0e-05
    stop: 5.0e-06
    steps: 10
lr_pen:
    function: stepwise
    start: ${scheduler.lr_net.start}
    stop: ${scheduler.lr_net.stop}
    steps: 10
dt:
    function: stepwise
    start: 0.02
    stop: 0.02
    steps: 10
chkpt_every_epoch: 300

Attributes:

Name	Type	Description
config	Namespace	Scheduler configuration.
scheduled_params	Namespace	Scheduled parameters.
network	Network	The neural network.
task	Task	The task being solved.
optimizer	Optimizer	The optimizer.
penalizer	Penalty	The penalty object.
stop_iter	int	Iteration to stop scheduling.
_current_iteration	int	Current iteration.

Source code in flyvision/solver.py

class HyperParamScheduler:
    """Schedules hyperparameters per training iteration.

    Calling the scheduler instance updates the respective hyperparameters per training
    iteration.

    Args:
        scheduler: Scheduler configuration.
        network: The neural network.
        task: The task being solved.
        optimizer: The optimizer.
        penalizer: The penalty object.

    Note: Default config in config/scheduler/scheduler.yaml
        ```yaml
        lr_net:
            function: stepwise
            start: 5.0e-05
            stop: 5.0e-06
            steps: 10
        lr_dec:
            function: stepwise
            start: 5.0e-05
            stop: 5.0e-06
            steps: 10
        lr_pen:
            function: stepwise
            start: ${scheduler.lr_net.start}
            stop: ${scheduler.lr_net.stop}
            steps: 10
        dt:
            function: stepwise
            start: 0.02
            stop: 0.02
            steps: 10
        chkpt_every_epoch: 300
        ```

    Attributes:
        config (Namespace): Scheduler configuration.
        scheduled_params (Namespace): Scheduled parameters.
        network (Network): The neural network.
        task (Task): The task being solved.
        optimizer (torch.optim.Optimizer): The optimizer.
        penalizer (Penalty): The penalty object.
        stop_iter (int): Iteration to stop scheduling.
        _current_iteration (int): Current iteration.
    """

    def __init__(
        self,
        scheduler: Namespace,
        network: Optional[Network],
        task: Optional[Task],
        optimizer: Optional[torch.optim.Optimizer],
        penalizer: Optional[Penalty],
    ):
        self.config = scheduler.deepcopy()
        self.scheduled_params = self.config.deepcopy()
        self.network = network
        self.task = task
        self.optimizer = optimizer
        self.penalizer = penalizer

        self.stop_iter = scheduler.get("sched_stop_iter", self.task.n_iters)
        self._current_iteration = 0

        self.scheduled_params = Namespace()
        for key, param in self.config.items():
            try:
                schedfn_config = SchedulerFunction(**param)
                logging.info("Init schedule for %s", key)
            except TypeError:
                # lazy way to skip the parameter if it's not a SchedulerFunction
                continue

            # these are the parameters that are scheduled
            param.array = getattr(self, schedfn_config.function)(
                self.stop_iter,
                self.task.n_iters,
                param.start,
                param.stop,
                param.steps,
            )
            self.scheduled_params[key] = param

    def __call__(self, iteration: int) -> None:
        """Update hyperparameters for the given iteration.

        Args:
            iteration: Current iteration.
        """
        self._current_iteration = iteration
        for key, param in self.scheduled_params.items():
            try:
                setattr(self, key, param.array[iteration])
            except IndexError as e:
                if iteration >= self.stop_iter:
                    setattr(self, key, param.array[-1])
                else:
                    raise e
        logging.info(self)

    def __repr__(self):
        return "Scheduler. Iteration: {}/{}.\nCurrent values: {}.".format(
            self._current_iteration,
            self.task.n_iters,
            self._params(),
        )

    def _params(self) -> dict:
        """Get current parameter values.

        Returns:
            A dictionary of current parameter values.
        """
        params = {}
        for key, _param in self.scheduled_params.items():
            value = getattr(self, key)
            params[key] = value
        return params

    # -------- Setter methods called automatically by the scheduler

    @property
    def dt(self) -> float:
        return self.task.dataset.dt

    @dt.setter
    def dt(self, value: float) -> None:
        self.task.dataset.dt = value

    @property
    def lr_net(self) -> float:
        if self.optimizer is None:
            return
        return self.optimizer.param_groups[0]["lr"]

    @lr_net.setter
    def lr_net(self, value: float) -> None:
        if self.optimizer is None:
            return
        self.optimizer.param_groups[0]["lr"] = value

    @property
    def lr_dec(self) -> list:
        if self.optimizer is None:
            return
        return [param_group["lr"] for param_group in self.optimizer.param_groups[1:]]

    @lr_dec.setter
    def lr_dec(self, value: float) -> None:
        if self.optimizer is None:
            return
        for param_group in self.optimizer.param_groups[1:]:
            param_group["lr"] = value

    @property
    def lr_pen(self) -> list:
        if self.penalizer is None:
            return
        return [
            param_group["lr"]
            for optim in self.penalizer.optimizers.values()
            for param_group in optim.param_groups
        ]

    @lr_pen.setter
    def lr_pen(self, value: float) -> None:
        if self.penalizer is None:
            return
        for optim in self.penalizer.optimizers.values():
            if optim is not None:
                for param_group in optim.param_groups:
                    param_group["lr"] = value

    @property
    def relu_leak(self) -> float:
        if self.network is None:
            return
        return getattr(self.network.dynamics.activation, "negative_slope", None)

    @relu_leak.setter
    def relu_leak(self, value: float) -> None:
        if self.network is None:
            return
        if hasattr(self.network.dynamics.activation, "negative_slope"):
            self.network.dynamics.activation.negative_slope = value

    @property
    def activity_penalty(self) -> float:
        if self.penalizer is None:
            return
        return self.penalizer.activity_penalty

    @activity_penalty.setter
    def activity_penalty(self, value: float) -> None:
        if self.penalizer is None:
            return
        self.penalizer.activity_penalty = value

    # -------- Decay Options

    @staticmethod
    def linear(
        stop_iter: int, n_iterations: int, start: float, stop: float, steps: int
    ) -> np.ndarray:
        """Generate a linear schedule from start to stop value."""
        f = np.linspace(start, stop, stop_iter)
        return np.pad(f, (0, n_iterations - len(f) + 1), constant_values=stop)

    @staticmethod
    def stepwise(
        stop_iter: int, n_iterations: int, start: float, stop: float, steps: int
    ) -> np.ndarray:
        """Generate a stepwise schedule from start to stop value."""
        f = np.linspace(start, stop, steps).repeat(stop_iter / steps)
        return np.pad(f, (0, n_iterations - len(f) + 1), constant_values=stop)

    @staticmethod
    def stepwise_2ndhalf(
        stop_iter: int, n_iterations: int, start: float, stop: float, steps: int
    ) -> np.ndarray:
        """Generate a stepwise schedule that decays in the second half of iterations."""
        f = np.linspace(start, stop, steps).repeat((stop_iter / 2) / steps)
        return np.pad(f, (n_iterations - len(f) + 1, 0), constant_values=start)

    @staticmethod
    def stepwise_half(
        stop_iter: int, n_iterations: int, start: float, stop: float, steps: int
    ) -> np.ndarray:
        """Generate a stepwise schedule that decays in the first half of iterations."""
        f = np.linspace(start, stop, steps).repeat((stop_iter / 2) / steps)
        return np.pad(f, (0, n_iterations - len(f) + 1), constant_values=stop)

    @staticmethod
    def steponential(
        stop_iter: int, n_iterations: int, start: float, stop: float, steps: int
    ) -> np.ndarray:
        """Generate an exponential stepwise schedule from start to stop value."""
        x = (1 / stop) ** (1 / steps)
        values = start / x ** np.arange(steps)
        f = values.repeat(stop_iter / steps)
        return np.pad(f, (0, n_iterations - len(f) + 1), constant_values=values[-1])

    @staticmethod
    def steponential_inv(
        stop_iter: int, n_iterations: int, start: float, stop: float, steps: int
    ) -> np.ndarray:
        """Generate an inverse exponential stepwise schedule."""
        _start = steps
        _stop = 0
        x = 1 / _stop
        values = _start / x ** np.arange(steps)
        f = values.repeat(stop_iter / steps)
        return np.pad(f, (0, n_iterations - len(f) + 1), constant_values=values[-1])

    @staticmethod
    def exponential(
        stop_iter: int, n_iterations: int, start: float, stop: float, steps: int
    ) -> np.ndarray:
        """Generate an exponential schedule from start to stop value."""
        tau = -stop_iter / (np.log(stop + 1e-15) - np.log(start))
        f = start * np.exp(-np.arange(stop_iter) / tau)
        return np.pad(f, (0, n_iterations - len(f) + 1), constant_values=stop)

    @staticmethod
    def exponential_half(
        stop_iter: int, n_iterations: int, start: float, stop: float, steps: int
    ) -> np.ndarray:
        """Generate exponential schedule that decays in the first half of iterations."""
        tau = -int((stop_iter / 2)) / (np.log(stop) - np.log(start))
        f = start * np.exp(-np.arange(int(stop_iter / 2)) / tau)
        return np.pad(f, (0, n_iterations - len(f) + 1), constant_values=stop)

__call__

__call__(iteration)

Update hyperparameters for the given iteration.

Parameters:

Name	Type	Description	Default
iteration	int	Current iteration.	required

Source code in flyvision/solver.py

def __call__(self, iteration: int) -> None:
    """Update hyperparameters for the given iteration.

    Args:
        iteration: Current iteration.
    """
    self._current_iteration = iteration
    for key, param in self.scheduled_params.items():
        try:
            setattr(self, key, param.array[iteration])
        except IndexError as e:
            if iteration >= self.stop_iter:
                setattr(self, key, param.array[-1])
            else:
                raise e
    logging.info(self)

linear staticmethod

linear(stop_iter, n_iterations, start, stop, steps)

Generate a linear schedule from start to stop value.

Source code in flyvision/solver.py

@staticmethod
def linear(
    stop_iter: int, n_iterations: int, start: float, stop: float, steps: int
) -> np.ndarray:
    """Generate a linear schedule from start to stop value."""
    f = np.linspace(start, stop, stop_iter)
    return np.pad(f, (0, n_iterations - len(f) + 1), constant_values=stop)

stepwise staticmethod

stepwise(stop_iter, n_iterations, start, stop, steps)

Generate a stepwise schedule from start to stop value.

Source code in flyvision/solver.py

@staticmethod
def stepwise(
    stop_iter: int, n_iterations: int, start: float, stop: float, steps: int
) -> np.ndarray:
    """Generate a stepwise schedule from start to stop value."""
    f = np.linspace(start, stop, steps).repeat(stop_iter / steps)
    return np.pad(f, (0, n_iterations - len(f) + 1), constant_values=stop)

stepwise_2ndhalf staticmethod

stepwise_2ndhalf(
    stop_iter, n_iterations, start, stop, steps
)

Generate a stepwise schedule that decays in the second half of iterations.

Source code in flyvision/solver.py

@staticmethod
def stepwise_2ndhalf(
    stop_iter: int, n_iterations: int, start: float, stop: float, steps: int
) -> np.ndarray:
    """Generate a stepwise schedule that decays in the second half of iterations."""
    f = np.linspace(start, stop, steps).repeat((stop_iter / 2) / steps)
    return np.pad(f, (n_iterations - len(f) + 1, 0), constant_values=start)

stepwise_half staticmethod

stepwise_half(stop_iter, n_iterations, start, stop, steps)

Generate a stepwise schedule that decays in the first half of iterations.

Source code in flyvision/solver.py

@staticmethod
def stepwise_half(
    stop_iter: int, n_iterations: int, start: float, stop: float, steps: int
) -> np.ndarray:
    """Generate a stepwise schedule that decays in the first half of iterations."""
    f = np.linspace(start, stop, steps).repeat((stop_iter / 2) / steps)
    return np.pad(f, (0, n_iterations - len(f) + 1), constant_values=stop)

steponential staticmethod

steponential(stop_iter, n_iterations, start, stop, steps)

Generate an exponential stepwise schedule from start to stop value.

Source code in flyvision/solver.py

@staticmethod
def steponential(
    stop_iter: int, n_iterations: int, start: float, stop: float, steps: int
) -> np.ndarray:
    """Generate an exponential stepwise schedule from start to stop value."""
    x = (1 / stop) ** (1 / steps)
    values = start / x ** np.arange(steps)
    f = values.repeat(stop_iter / steps)
    return np.pad(f, (0, n_iterations - len(f) + 1), constant_values=values[-1])

steponential_inv staticmethod

steponential_inv(
    stop_iter, n_iterations, start, stop, steps
)

Generate an inverse exponential stepwise schedule.

Source code in flyvision/solver.py

@staticmethod
def steponential_inv(
    stop_iter: int, n_iterations: int, start: float, stop: float, steps: int
) -> np.ndarray:
    """Generate an inverse exponential stepwise schedule."""
    _start = steps
    _stop = 0
    x = 1 / _stop
    values = _start / x ** np.arange(steps)
    f = values.repeat(stop_iter / steps)
    return np.pad(f, (0, n_iterations - len(f) + 1), constant_values=values[-1])

exponential staticmethod

exponential(stop_iter, n_iterations, start, stop, steps)

Generate an exponential schedule from start to stop value.

Source code in flyvision/solver.py

@staticmethod
def exponential(
    stop_iter: int, n_iterations: int, start: float, stop: float, steps: int
) -> np.ndarray:
    """Generate an exponential schedule from start to stop value."""
    tau = -stop_iter / (np.log(stop + 1e-15) - np.log(start))
    f = start * np.exp(-np.arange(stop_iter) / tau)
    return np.pad(f, (0, n_iterations - len(f) + 1), constant_values=stop)

exponential_half staticmethod

exponential_half(
    stop_iter, n_iterations, start, stop, steps
)

Generate exponential schedule that decays in the first half of iterations.

Source code in flyvision/solver.py

@staticmethod
def exponential_half(
    stop_iter: int, n_iterations: int, start: float, stop: float, steps: int
) -> np.ndarray:
    """Generate exponential schedule that decays in the first half of iterations."""
    tau = -int((stop_iter / 2)) / (np.log(stop) - np.log(start))
    f = start * np.exp(-np.arange(int(stop_iter / 2)) / tau)
    return np.pad(f, (0, n_iterations - len(f) + 1), constant_values=stop)

SchedulerFunction dataclass

Configuration for a scheduler function.

Attributes:

Name	Type	Description
start	float	Start value.
stop	float	Stop value.
steps	int	Number of steps.
function	str	Name of the scheduling function.

Source code in flyvision/solver.py

@dataclass
class SchedulerFunction:
    """Configuration for a scheduler function.

    Attributes:
        start (float): Start value.
        stop (float): Stop value.
        steps (int): Number of steps.
        function (str): Name of the scheduling function.
    """

    start: float
    stop: float
    steps: int
    function: str