OpenLTH: A Framework for Lottery Tickets and Beyond
This framework implements key experiments from recent work on the lottery ticket hypothesis and the science of deep learning:
- The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks. Jonathan Frankle & Michael Carbin. ICLR 2019.
- Stabilizing the LTH/The LTH at Scale. Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M. Roy, & Michael Carbin. Arxiv.
- Linear Mode Connectivity and the LTH. Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M. Roy, & Michael Carbin. Arxiv.
- The Early Phase of Neural Network Training. Jonathan Frankle, David J. Schwab, and Ari S. Morcos. ICLR 2020.
- Training BatchNorm and Only BatchNorm. Jonathan Frankle, David J. Schwab, & Ari S. Morcos. Arxiv
It was created by Jonathan Frankle during his time as a summer intern and student researcher at FAIR starting in June 2019. It is his current working research codebase for image classification experiments on the lottery ticket hypothesis.
If you use this library in a research paper, please cite this repository.
OpenLTH is licensed under the MIT license, as found in the LICENSE file.
We welcome your contributions! See the CONTRIBUTING file for how to help out.
Table of Contents
This framework is designed with four goals in mind:
- To train standard neural networks for image classification tasks.
- To run pruning and lottery ticket experiments.
- To automatically manage experiments and results without any manual intervention.
- To make it easy to add new datasets, models, and other modifications.
1.2 Why Another Framework?
This framework and its predecessors were developed in the course of conducting research on the lottery ticket hypothesis.
- Pruning. Those experiments involve pruning neural networks. Pruning is a first-class citizen of the framework.
- Hyperparameter management. Those experiments involve extensive hyperparameter sweeps. Default hyperparameters are easy to modify, and results are automatically indexed by the hyperparameters (so you never need to worry about naming experiments).
- Extensibility. Those experiments study an ever-growing range of models and datasets. Consistent abstractions for models and datasets make it easy to add new ones in a modular way.
- Enabling new experiments. Those experiments involve adding new hyperparameters. Each hyperparameter automatically surfaces on the command-line and integrates into experiment naming in a backwards-compatible way.
- Re-training and ablation. Those experiments rely on many post-hoc re-training and ablation experiments. Such experiments can be created by writing a single function, after which they are automatically available to be called from the command line and stored in a standard way.
- Platform flexibility. Those experiments have had to run on many different platforms and providers. The notion of a "platform" is a first-class abstraction, making it easy to adapt the framework to new settings or to multiple settings.
- Python 3.7 or greater (this framework extensively uses features from Python 3.7)
- PyTorch 1.4 or greater
- TorchVision 0.5.0 or greater
- NVIDIA Apex (optional, but required for 16-bit training)
- Install the requirements.
- Clone this repository.
platforms/local.pyso that it contains the paths where you want datasets and results to be stored. By default, they will be stored in
~/open_lth_datasets/. To train with ImageNet, you will need to specify the path where ImageNet is stored.
2.3 The Command-Line Interface
All interactions with the framework occur through the command-line interface:
In response, you will see the following message.
================================================================================== OpenLTH: A Framework for Research on Lottery Tickets and Beyond ---------------------------------------------------------------------------------- Choose a command to run: * open_lth.py train [...] => Train a model. * open_lth.py lottery [...] => Run a lottery ticket hypothesis experiment. * open_lth.py lottery_branch [...] => Run a lottery branch. ==================================================================================
This framework has three subcommands for its three experimental workflows:
train (for training a network),
lottery (for running a lottery ticket hypothesis experiment), and
lottery_branch (for running an ablation on a lottery ticket experiment). To learn about adding more experimental workflows, see
2.4 Training a Network
To train a network, use the
train subcommand. You will need to specify the model to be trained, the dataset on which to train it, and other standard hyperparameters (e.g., batch size, learning rate, training steps). There are two ways to do so:
- Specify these hyperparameters by hand. Each hyperparameter is controlled by a command-line argument. To see the complete list, run
python open_lth.py train --help. Many hyperparameters are required (e.g.,
--lr). Others are optional (e.g.,
- Use the defaults. Each model comes with a set of defaults that achieve standard performance. You can specify the name of the model you wish to train using the
--default_hparamsargument and load the default hyperparameters for that model. You can still override any default using the individual arguments for each hyperparameter.
In practice, you will almost always begin with a set of defaults and optionally modify individual hyperparameters as desired. To view the default hyperparameters for ResNet-20 on CIFAR-10, use the following command. (For a full list of available models, see 2.11.) Each of the hyperparameters from before will be updated with its default value.
python open_lth.py train --default_hparams=cifar_resnet_20 --help
To train with these default hyperparameters, use the following command (that is, leave off
python open_lth.py train --default_hparams=cifar_resnet_20
The training process will then begin. The framework will print the required and non-default hyperparameters for the training run and the location where the resulting model will be stored.
================================================================================== Training a Model (Replicate 1) ---------------------------------------------------------------------------------- Dataset Hyperparameters * dataset_name => cifar10 * batch_size => 128 Model Hyperparameters * model_name => cifar_resnet_20 * model_init => kaiming_normal * batchnorm_init => uniform Training Hyperparameters * optimizer_name => sgd * lr => 0.1 * training_steps => 160ep * momentum => 0.9 * milestone_steps => 80ep,120ep * gamma => 0.1 * weight_decay => 0.0001 Output Location: /home/jfrankle/open_lth_data/train_71bc92a970b64a76d7ab7681764b0021/replicate_1/main ==================================================================================
Before each epoch, it will print the test error and loss.
test ep 000 it 000 loss 68.261 acc 10.00% time 0.00s
To suppress these messages, use the
--quiet command-line argument.
To override any default hyperparameter values, use the corresponding hyperparameter arguments. For example, to increase the batch size and learning rate and add 10 epochs of learning rate warmup:
python open_lth.py train --default_hparams=cifar_resnet_20 --batch_size=1024 --lr=0.8 --warmup_steps=10ep
2.5 Running a Lottery Ticket Experiment
A lottery ticket experiment involves repeatedly training the network to completion, pruning weights, rewinding unpruned weights to their value at initialization, and retraining. For more details, see The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks.
To run a lottery ticket experiment, use the
lottery subcommand. You will need to specify all the same hyperparameters required for training. In addition, you will need to specify hyperparameters for pruning the network. To see the complete set of hyperparameters, run:
python open_lth.py lottery --help
For pruning, you will need to specify a value for the
--pruning_strategy hyperparameter. By default, the framework includes only one pruning strategy: pruning the lowest magnitude weights globally in a sparse fashion (
sparse_global). (For instructions on adding new pruning strategies, see Section 4.7.)
Once again, it is easiest to load the default hyperparameters for a model (which includes the pruning strategy and other pruning details) using the
--default_hparams argument. In addition, you will need to specify the number of times that the network should be pruned, rewound, and retrained, known as the number of pruning levels. To do so, use the
--levels flag. Level 0 is training the full network; specifying
--levels=3 would then prune, rewind, and retrain a further three times.
To run a lottery ticket experiment with the default hyperparameters for ResNet-20 on CIFAR-10 with three pruning levels, use the following command:
python open_lth.py lottery --default_hparams=cifar_resnet_20 --levels=3
2.6 Running a Lottery Ticket Experiment with Rewinding
In the original paper on the lottery ticket hypothesis, unpruned weights were always rewound to their values from initialization before the retraining phase. In recent work (Stabilizing the LTH/The LTH at Scale, Linear Mode Connectivity and the LTH and The Early Phase of Neural Network Training), weights are typically rewound to their values from step
k during training (rather than from initialization).
This framework incorporates that concept through a broader feature called pretraining. Optionally, the full network can be pretrained for
k steps, with the resulting weights used as the starting point for the lottery ticket procedure. Rewinding to step
k and pretraining for
k steps are functionally identical, but pretraining offers increased flexibility. For example, you can pretrain using a different training set, batch size, or loss function; this is precisely the experiment performed in Section 5 of The Early Phase of Neural Network Training.
If you wish to use the same hyperparameters for pretraining and the training process itself (i.e., perform standard lottery ticket rewinding), you can set the argument
--rewinding_steps. For example, to rewind to iteration 500 after pruning (or, equivalently, to pretrain for 500 iterations), use the following command:
python open_lth.py lottery --default_hparams=cifar_resnet_20 --levels=3 --rewinding_steps=500it
If you wish to have different behavior during the pre-training phase (e.g., to pre-train with self-supervised rotation or even on a different task), use the
--pretrain argument. After doing so, the
--help interface will offer the full suite of pretraining hyperparameters, including learning rate, batch size, dataset, etc. By default, it will pretrain using the same hyperparameters as specified for standard training. You will noueed to set the
--pretrain_training_steps argument to the number of steps for which you wish to pretrain. Note that the network will still only train for the number of steps specified in
--training_steps. Any steps specified in
--pretrain_training_steps will be subtracted from
--training_steps. In addition, the main phase of training will start from step
--pretrain_training_steps, including the learning rate and state of the dataset at that step.
2.7 Lottery Ticket Branches
Many of the experiments in the lottery ticket papers involve running ablations based on the main lottery ticket experiment. For example, training the same pruned network but with a different random initialization or training with the same random initialization but a random pruning mask. Typically, you run this ablation for the networks produced by the main lottery ticket experiment at each level of pruning. For example, if we run a lottery ticket experiment with three levels of pruning, we will run this ablation four times - once on the unpruned network and once on the network at each level of pruning. We refer to such experiments as branches, since they branch off of the main lottery ticket trunk.
To run a branch, you need to first run the corresponding lottery ticket trunk. Afterwards, use the
lottery_branch subcommand. This subcommand must be followed by the name of the branch (a sub-sub command).
python open_lth.py lottery_branch
If you do not specify a branch name, the framework will list all available branches.
================================================================================== OpenLTH: A Framework for Research on Lottery Tickets and Beyond ---------------------------------------------------------------------------------- Choose a branch to run: * open_lth.py lottery_branch randomly_prune [...] => Randomly prune the model. * open_lth.py lottery_branch randomly_reinitialize [...] => Randomly reinitialize the model. * open_lth.py lottery_branch retrain [...] => Retrain the model with different hyperparameters. ==================================================================================
Each branch takes the standard lottery ticket arguments (in order to specify the trunk off of which to branch) and several branch-specific arguments. To see these arguments, use the
--help command. Many of these arguments have standard default values, so you may not need to specify some or all of them. Finally, each
lottery_branch command runs for a range of pruning levels (for example, '0-2,5-8,12'), so you need to specify that level with the
For example, to randomly reinitialize level 2 of the prior lottery ticket experiment with random seed
7, use the following command:
python open_lth.py lottery_branch randomly_prune --default_hparams=cifar_resnet_20 --levels=0-20 --seed=7
2.8 Accessing Results
All experiments are automatically named according to their hyperparameters. Specifically, all required hyperparameters and all optional hyperparameters that are specified are combined in a canonical fashion and hashed. This hash is the name under which the experiment is stored. The results of a training run are then stored under:
<root> is the data root directory stored in
platforms/local.py; it defaults to
The results themselves are stored in a file called
logger, which only appears after training is complete. This file is a lightweight CSV where each line is one piece of telemetry data about the model. A line consists of three comma-separated values: the name of the kind of telemetry (e.g.,
test-accuracy), the iteration of training at which the telemetry data was collected (e.g.,
391), and the value itself. You can parse this file manually, or use
training/metric_logger.py, which is used by the framework to read and write these files.
To get the name of the output location for a particular run, use the
python open_lth.py train --default_hparams=cifar_resnet_20 --display_output_location /home/jfrankle/open_lth_data/train_71bc92a970b64a76d7ab7681764b0021/replicate_1/main
2.9 Checkpointing and Re-running
Each experiment will automatically checkpoint after every epoch. If you re-launch a job, it will automatically pick up from there. If a job has already completed, it will not run again unless you manually delete the associated results.
If you wish to run multiple copies of an experiment (which is good scientific practice), use the
--replicate argument. This optional argument specifies the replicate number of an experiment. For example,
--replicate=5 will store the experiment under
If no replicate is specified,
--replicate will default to 1.
2.10 Modifying the Training Environment
To suppress the outputs, use the
To specify the number of PyTorch worker threads used to load data, use the
--num_workers argument. This value defaults to 0, although you will need to specify a value for training with ImageNet.
The framework will automatically use all available GPUs. To change this behavior, you will need to modify the number of visible GPUs using the
CUDA_VISIBLE_DEVICES environment variable.
2.11 Available Models
The framework includes the following models. Each model family shares the same default hyperparameters.
|Model||Description||Name in Framework||Example|
|LeNet||A fully-connected MLP.||
|VGG for CIFAR-10||A convolutional network with max pooling in the style of VGG.||
|ResNet for CIFAR-10||A residual network for CIFAR-10. This is a different family of architectures than those designed for ImageNet.||
|Wide ResNet for CIFAR-10||The ResNets for CIFAR-10 in which the width of the network can be varied.||
|ResNet for ImageNet||A residual network for ImageNet. This is a different family of architectures than those designed for CIFAR-10, although they can be trained on CIFAR-10.||
|Wide ResNet for ImageNet||The ResNets for ImageNet in which the width of the network can be varied.||
This framework includes standard ResNet models for ImageNet and a standard data preprocessing for ImageNet. Using the default hyperparameters and 16-bit precision training,
imagenet_resnet_50 trains to 76.1% top-1 accuracy in 22 hours on four V100-16GB GPUs. To use ImageNet, you will have to take additional steps.
- Prepare the ImageNet dataset.
- Create two folders,
val, each of which has one subfolder for each class containing the JPEG images of the examples in that class.
platforms/local.pyto return this location.
- Create two folders,
- If you wish to train with 16-bit precision, you will need to install the NVIDIA Apex and add the
--apex_fp16argument to the training command.
This framework is designed to be extensible, making it easy to add new datasets, models, initializers, optimizers, pruning strategies, hyperparameters, branches, workflows, and other customizations. This section discusses the internals. Section 4 is a how-to guide for extending the framework.
Note that this framework makes extensive use of Python Data Classes, a feature introduced in Python 3.7. You will need to understand this feature before you dive into the code. This framework also makes extensive use of object oriented subclassing with the help of the Python ABC library.
3.1 Hyperparameters Abstraction
The lowest-level abstraction in the framework is an object that stores a bundle of hyperparameters. The abstract base class for all such bundles of hyperparameters is the
Hparams Data Class, which can be found in
foundations/hparams.py. This file also includes four subclasses of
Hparams that are used extensively throughout the framework:
DatasetHparams(which includes all hyperparameters necessary to specify a dataset, like its name and batch size)
ModelHparams(which includes all hyperparameters necessary to specify a model, like its name and initializer)
TrainingHparams(which includes all hyperparameters necessary to describe how to train a model, like the optimizer, learning rate, warmup, annealing, and number of training steps)
PruningHparams(which is the base class for the hyperparameters required by each pruning strategy)
Each field of these dataclasses is the name of the hyperparameter. The type annotation is the type that the hyperparameter must have. If a default value is specified for the field, that is the default value for the hyperparameter; if no default is provided, then the hyperparameter is required and must be specified manually.
Hparams subclass must also set the
_description fields with default values that describe the nature of this bundle of hyperparameters. It may optionally include a string field
_hyperparameter with a default string value that describes the hyperparameter and how it should be set. For example, in addition to the
TrainingHparams has the
_lr field that explains how the
lr field should be set.
Hparams subclass provides several behaviors to its subclasses. Most importantly, it has a static method
add_args which takes as input a Python command-line
ArgumentParser and adds each of the hyperparameters as a flag
--hyperparameter. Since each hyperparameter has a name, type annotation, and possibly a default value and help text (the
_hyperparameter field), it can be converted into a command-line argument automatically. This is how the per-hyperparameter command-line arguments are populated. This function optionally takes an instance of the class that overrides default values; this is how
--default_hparams is implemented. Corresponding to the
add_args static method is a
create_from_args static method that creates an instance of the class from a Python
argparse.NameSpace object that results from using the
Hparams object has a
__str__ method that converts an instance into a string in a canonical way. During this conversion, any hyperparameters that are set to their default values are left off. This step is very important for ensuring that models are saved in a backwards compatible way as new hyperparameters are added.
3.2 Modules for Datasets, Models, Training, and Pruning
Running a lottery ticket experiment involves combining four largely independent components:
- There must be a way to retrieve a dataset.
- There must be a way to retrieve a model.
- There must be a way to train a model on a dataset.
- Three must be a way to prune a model.
This framework breaks these components into distinct modules that are as independent as possible. The common specification for these modules is the
Hparams objects. To request a dataset from the dataset module, provide a
DatasetHparams instance. To request a model from the models module, provide a
ModelHparams instance. To train a model, provide a dataset, a model, and a
TrainingHparams instance. To prune a model, provide the model and a
PruningHparams object. The inner workings of these modules can be understood largely independently from each other, with a few final abstractions to glue everything together.
3.3 The Datasets Module
Each dataset consists of two abstractions:
Datasetthat stores the dataset, labels, and any data augmentation.
DataLoaderthat loads the dataset for training or testing. It must keep track of the batch size, multithreaded infrastructure for data loading, and random shuffling.
A dataset must subclass the
DataLoader abstract base classes in
datasets/base.py. Both of these classes subclass the corresponding PyTorch
DataLoader classes, although they have a richer API to facilitate functionality in other modules and to enable build-in transformations like subsampling, random labels, and blurring.
For simple datasets that can fit in memory, these base classes provide most of the necessary functionality, so the subclasses are small. In fact, MNIST (
datasets/mnist.py) and CIFAR-10 (
datasets/cifar10.py) use the base
DataLoader without modification. In contrast, ImageNet (
datasets/imagenet.py) replaces all functionality due to the specialized needs of loading such a large dataset efficiently.
The external interface of this module is contained in
datasets/registry.py. The registry contains a list of all existing datasets in the framework (so that they can be discovered and loaded). Its most important function is
get(), which takes as input a
DatasetHparams instance and a boolean specifying whether to load the train or test set; it returns the
DataLoader object corresponding to the
DatasetHparams (i.e., with the right batch size and additional transformations). This module also contains a function for getting the number of
iterations_per_epoch() and the
num_classes() corresponding to a particular
DatasetHparams, both of which are important for other modules.
3.4 The Models Module
Each model is created by subclassing the
Model abstract base class in
models/base.py. This base class is a valid PyTorch
Module with several additional abstract methods that support other functionality throughout the framework. In particular, any subclass must have static methods to determine whether a string model name (e.g.,
cifar_resnet_20) is valid and to create a model object from a string name, a number of outputs, and an initializer.
It must also have instance properties that return the names of the tensors in the output layer (
output_layer_names) and all tensors that are available for pruning (
prunable_layer_names - by default just the kernels of convolutional and linear layers). These properties are used elsewhere in the framework for transfer learning, weight freezing, and pruning.
Finally, it must have a static method that returns the set of default hyperparameters for the corresponding model family (as
Hparams objects); doing so makes it possible to load the default hyperparameters rather than specifying them one by one on the command line.
Otherwise, these models are identical to standard
Modules in PyTorch.
The external interface of this module is contained in
datasets/registry.py, there is a list of all existing models in the framework so that they can be discovered and loaded. The registry similarly contains a
get() function that, given a
ModelHparams instance and the number of outputs, returns the corresponding
Model as specified. In the course of creating a model, the
get() function also loads the initializer and BatchNorm initializer specified in the
ModelHparams instance. All initializers are functions stored in
models/initializers.py, and no registration is necessary. All BatchNorm initializers are functions stored in
models/bn_initializers.py, and no registration is necessary.
Finally, the registry has functions for getting the default hyperparameters for a model, loading a saved model from a path, and checking whether a path contains a saved model.
3.5 The Step Abstraction
In several places throughout the framework, it is necessary to keep track of a particular "step" of training. Depending on the particular framework, a step takes one of two forms: an iteration of training or an epoch number and an iteration offset into that epoch. In some places in this framework, it is easier to use one representation than the other. To make it easy to convert back and forth between these representations, all steps are stored as
Step objects (
foundations/step.py). A step object can be created from either representation, but it requires the number of iterations per epoch so that it can convert back and forth between these two representations.
3.6 The Training Module
The training module centers on a single function: the
train() function in
training/train.py. This function takes a
Model and a
DataLoader as arguments along with a
TrainingHparams instance. It then trains the
Model on the dataset provided by the
DataLoader as specified in the training hyperparameters.
train() function takes four other arguments. These include the
output_location where all results should be stored, the optional
Step object at which training should begin (the learning rate schedule and training set are advanced to this point), and the optional
Step object at which training should end (if not the default value specified in the
Most importantly, it takes an argument called
callbacks. This argument requires some explaining. A key design goal of the training module is to keep the main training loop in
train() as clean as possible. This means that the loop should only contain standard setup, training, checkpointing behavior, and a
MetricLogger to record telemetry. The loop should not be modified to add other behaviors, like saving the network state, running the test set, adding if-statements for new experiments, etc.
Instead, the behavior of the loop is modified by providing callback functions (known as hooks in other frameworks). These callbacks are called before every training step and after the final training step. They are provided with the current training step, the model, the optimizer, the output location, and the logger, and they can perform functions like saving the model state, running the test set, checkpointing, printing useful debugging information, etc. As new functionality is needed in the training loop, simply create new callbacks.
training/standard_callbacks.py contains a set of the most common callbacks you are likely to use, like evaluating the model on a
DataLoader or saving the model state. It also contains a set of higher-order functions that modify a callback to run at a certain step or interval. Finally, it includes a set of standard callbacks for a training run:
- Save the network before and after training
- Run the test set every epoch
- Update the checkpoint every epoch
- Save the logger after training
training/train.py contains a function called
standard_train() that takes a model, dataset and training hyperparameters, and an output location as inputs and trains the model using the standard callbacks and the main training loop. This function is used by the
To create optimizers and learning rate scheduler objects,
train() calls the
get_lr_schedule() functions in
training/optimizers.py, which serve as small-scale registries for these objects.
3.7 The Pruning Module
The pruning module is in the
pruning/ directory. It contains a
Mask abstraction, which keeps track of the binary mask
Tensor for each prunable layer in a model. The module keeps track of different pruning strategy classes (subclasses of the abstract base class
pruning/base.py. Each pruning strategy has two members:
- A static method
get_pruning_hparams()that returns a subclass (not an instance) of the
foundations/hparams.py. Since different pruning methods may require different hyperparameters, each pruning method is permitted to specify its own
PruningHparamsobject. This object is used to generate the command-line arguments for the pruning strategy specified by the
- A static method
prunethat takes a
PruningHparamsinstance, a trained
Model, and the current
Maskand returns a new mask representing one further pruning step according to the hyperparameters.
The external interface of this module is contained in
pruning/registry.py. Like the other registries, it has a
get() function for getting a pruning class from a
PruningHparams instance. It also has a
get_pruning_hparams function for getting the
PruningHparams subclass for a particular pruning strategy.
Finally, this module contains a
PrunedModel class (in
pruning/pruned_model.py). This class is a wrapper around a
models/base.py) that applies a
Mask object to prune weights. This class is used heavily by the lottery ticket and branch experiments to effectuate pruning.
These individual components (datasets, models, training, and pruning) come together to allow for training workflows. The framework currently has two training workflows: training a model normally (the
train subcommand) and running a lottery ticket experiment (the
Each of these workflows requires a slightly different set of hyperparameters. Training a model requires
TrainingHparams instances (but notably no
PruningHparams, since no pruning occurs in this workflow). In contrast, a lottery ticket experiment also needs
PruningHparams and, optionally, a separate set of
TrainingHparams for pre-training.
In summary, each workflow needs a "bundle" of
Hparams objects of different kinds. The framework represents this abstraction with a descriptor object, which describes everything necessary to conduct the workflow (training a network or running a lottery ticket experiment). These objects descend from
foundations/desc.py, which contains the abstract base class
Desc. This class is a Python dataclass whose fields are
Hparams objects. It requires subclasses to implement a
create_from_args static methods that create the necessary command-line arguments like the similar methods in the
Hparams base class; typically, these functions will simply call the corresponding ones in the constituent
Desc base class contains the function that makes automatic experiment naming possible. It has a property called
hashname that combines all
Hparams objects in its fields into a single string in a canonical way and returns the MD5 hash. This hash later becomes the name under which each experiment is stored. It is therefore important to be careful when modifying
Desc objects, as doing so may break backwards compatibility with the hashes of pre-existing experiments.
The training and lottery workflows contain subclasses of
lottery/desc.py. Each of these subclasses contains the requisite fields and implements the required abstract methods. They also include other useful properties derived from their constituent hyperparameters for the convenience of higher-level abstractions.
3.9 Wiring Everything Together with Runners
A descriptor has everything necessary to specify how a particular network should be trained, but it is missing other meta-data necessary to fully describe a run. For example, a training run needs a replicate number, and a lottery run needs to know the number of pruning levels for which to run. This information is captured in a higher-level abstraction known as a
Each runner (subclasses of
foundations/runner.py) has static
create_from_args methods that interact with command-line arguments, calling the same methods on their requisite descriptors and adding other runner-level arguments. Once a
Runner instance has been created, the
run() method (which takes no arguments) initiates the run. This includes creating a model and making one or more calls to
train() depending on the details of the runner. For example, the runner for the
train subcommand (found in
training/runner.py) performs a single training run; the runner for the
lottery subcommand (found in
lottery/runner.py) pretrains a network, trains it, prunes it, and then repeatedly re-trains and prunes it using the
PrunedModel class and a pruning
The runners are the highest level of abstraction, connecting directly to the command-line interface. Each runner must be registered in
cli/runner_registry.py. The name under which it is added is the name of the subcommand used to access it on the command-line.
A lottery ticket training run with
N pruning levels produces
N+1 networks: the original, unpruned network and the networks that result from each of the
N pruning levels. In many lottery ticket experiments, some subsequent ablation is conducted on each of these networks. For example, they are randomly reinitialized and retrained to assess the role of the original initialization or they are randomly pruned to assess the role of the specific sparsity pattern found by pruning.
We refer to such experiments as branches, since they branch off of the main lottery ticket trunk. Branches are implemented in the
lottery/branch directory. Each branch is implemented by writing a subclass of the
Branch abstract base class in
lottery/branch/base.py. This subclass must have three methods:
- A static method
namethat returns the name of the branch as is used to refer to it on the command line.
- A static method
descriptionthat returns a description of the branch that is used on the command line.
- An instance method
branch_functionthat contains the body of the branch. This function receives all of the state of the superclass instance, including a
LotteryDescobject describing the run being branched, the replicate being branched, and the level being branched. The
branch_functionmethod can have arguments that vary its behavior; these arguments are automatically converted into command-line arguments when this branch is called. The
branch_rootproperty of the superclass instance contains the name of the output directory that should be used for a call to this branch; it is generated based on the arguments with which the function is called.
The only remaining step is to register this branch in
The branch module is designed to make it as easy as possible to write a new branch, since this is a heavily-used path in lottery ticket research. All that it takes to write a branch is to write a
branch_function; the base class takes care of connecting it to the command line and initializing its state automatically.
This magic relies on a little meta-programming to work. The base
Branch class in
lottery/branch/base.py uses the
__init_subclass__ function to modify any subclass that is created. It inspects the signature of the
branch_function subclass, extracting the name, type annotation (which is required), and default value of each argument. Together, this information is enough to dynamically generate a
Hparams class in which these arguments become hyperparameters.
Using that dynamically generated
__init_subclass__ dynamically generates a descriptor for the subclass using the
make_BranchDesc function in
lottery/branch/desc.py. This descriptor has one set of
Hparams (those that were dynamically generated) and a
LotteryDesc field containing the descriptor of the lottery ticket trunk that it is building on.
Branch base class itself is also a
Runner; it builds the descriptor and, when run, calls
branch_function in the subclass, executing the branch.
To connect the branch to the command line, there is another runner (in
lottery/branch/runner.py). When the
lottery_branch subcommand is used, it requests the name of the desired branch as a sub-subcommand. Once that is provided, it consults the branch registry (
lottery/branch/registry.py) and dispatches to the appropriate branch
Runner to execute the job.
It is typical to use the same codebase on many different infrastructures (such as a local machine, a cluster, and one or more cloud providers). Each of these infrastructures will have different locations where results and datasets will be stored and different ways of accessing filesystems. They may even need to call the runner's
run() functions in a different fashion.
To make it easy to run this framework on multiple infrastructures, it includes a
Platform abstraction. Each
Platform class describes where to find resources (like datasets), where to store results, what hardware is available (if there are GPUs and how many if so) and how to run a job on the platform. Arguments may be required to create a
Platform instance, for example the number of worker threads to use.
To enable this behavior, each
Platform object is a dataclass that descends from
Hparams; this makes it possible for its fields to be converted into command-line arguments and for an instance to be created from these arguments. The abstract base
Platform class that all others subclass (found in
platforms/base.py) contains a field for the number of worker threads to use for data loading.
It also has abstract properties that specify where data should be found and results stored; these must be implemented by each subclass.
Finally, it has a series of static methods that mediate access to the filesystem; by default, these are set to use the standard Python commands for the local filesystem, but it may be important to override them on certain infrastructures.
Finally, it has a method called
run_job() that receives a function
f as an argument, performs any pre-job setup, and calls the function. Most importantly, this function installs the
Platform instance as the global platform for the entire codebase. In practice, this entails modifying the global variable
platforms/platform.py. Throughout the codebase, modules look to this global variable (accessed through the
get_platform() function in
platforms/platform) to determine where data is stored, the hardware on which to run a job, etc. It was cleaner to make the current platform instance a global rather than to carry it along through every function call in the codebase.
local platform will automatically use all GPUs available using PyTorch
DataParallel. If you choose to do distributed training, the
base platform includes primitives for distributed training like
barrier; the codebase calls all of these functions in the proper places so that it is forward-compatible with distributed training should you choose to use it.
All platform subclasses must be registered in
platforms/registry.py, which makes them available for use at the command line using the
--platform argument. By default the
local platform (which runs on the local machine) is used.
This codebase contains extensive unit tests for the low-level modules and the lottery ticket hypothesis runner pipeline. To execute these tests, run the following command:
python -m unittest discover
The unit tests include a regression test for the directory names generated by the framework to ensure that new hyperparameters have not inadvertently changed existing names. Note that the unit tests do not directly test the command-line interface, the train
Runner, and the branch infrastructure.
Please read Section 3 before trying to extend the framework. Careless changes can have unexpected consequences, such as breaking backwards compatibility and making it impossible for the framework to access your existing models.
4.1 Adding a New Dataset
Create a new file in the
datasets directory that subclasses the abstract base classes
datasets/base.py with classes that are also called
datasets/registry.py to import this module and add the module (not the classes in the module) to the dictionary of
registered_datasets with the name that you wish for it to be called. For small datasets that fit in memory (e.g., SVHN), use
datasets/cifar10.py as a template and take advantage of functionality built into the base classes. For larger datasets (e.g., Places), use
datasets/imagenet.py as a template; you may need to throw away functionality in the base classes.
4.2 Adding a New Model
Create a new file in the
models directory that subclasses the abstract base class
models/registry.py to import this module and add the class (not the module containing the class) to the list of
registered_models. As a template, use
4.3 Adding a New Initializer
Add the new initializer function to
models/initializers.py under the name that you want it to be called. To add a new BatchNorm initializer, do the same in
models/bn_initializers.py. No registration is necessary in either case.
4.4 Adding a New Optimizer
Modify the if-statement in the
get_optimizer function of
training/optimizers.py to create the new optimizer when the appropriate hyperparameters are specified.
4.5 Adding a New Hyperparameter
Modify the appropriate set of hyperparameters in
foundations/hparams.py to include the desired hyperparameter. The hyperparameter must have a default value, and this default value must eliminate the effect of the hyperparameter. The goal is to ensure that adding this hyperparameter is backwards compatible. This default value should ensure that all preexisting models would train in the same way if this hyperparameter had been present and set to its default value.
If the new hyperparameter doesn't have a default value, then it will change the way results directory names are computed for all preexisting models, making it impossible for the framework to find them. If the default value is not a no-op, then all preexisting models (where were trained under the implicit assumption that this hyperparameter was set to its default value) will no longer be valid.
The unit tests include a regression test for the directory names generated by the framework to ensure that new hyperparameters have not inadvertently changed existing names. Be sure to run the unit tests after adding a new hyperparameter.
4.6 Modifying the Training Loop
Where possible, try to modify the training loop by creating a new kind of optimizer, a new kind of loss function, or a new callback. New callbacks can be added to
training/train.py, gated by a new hyperparameter. The training loop is designed to be as clean and pared down as possible and to use callbacks and the other objects to abstract away the complicated parts, so try to avoid modifying the loop if at all possible. If you need to access the gradients, consider adding a second set of
post_gradient_callbacks that are called after the gradients are computed but before the optimizer steps. This would be a new argument for
train() and possibly
4.7 Adding a New Pruning Strategy
Create a new file in the
pruning directory that subclasses the abstract base class
pruning/base.py. The new pruning strategy needs a static method that returns the hyperparameters it requires (recall that each pruning method can have a different set of hyperparameters). Modify
pruning/registry.py to import this module and add the class (not the module containing the class) to the dictionary of
registered_strategies under the key that you want to use to describe this strategy going forward.
4.8 Adding a New Workflow
- Create a new directory to store the workflow.
- Create a file with a descriptor data class that subclasses from
foundations/desc.py; it should have fields for any
Hparamsobjects necessary to describe the workflow. It should implement the
name_prefixstatic methods as necessary for the desired behavior.
- Create a file with a runner class that subclasses from
foundations/runner.py. Create a constructor or make the runner a data class. Implement the
create_from_argsstatic methods to interact with the command line. Implement the
descriptionstatic method to describe the runner. Implement the
display_output_locationinstance method to respond to the
--display_output_locationcommand-line argument. Finally, create the
runinstance method with the logic for performing any training necessary for the workflow.
- Register the runner in
4.9 Adding a New Branch
Create a new file in the
lottery/branch directory that subclasses the abstract base class
lottery/branch/base.py. This subclass needs three functions:
- An instance method called
branch_function. This function executes the branch, including any training that is required. The first argument must be
self(the instance). Give the function any additional arguments that you want to appear as command-line arguments when this branch is run. These arguments must have type annotations, and they may have default values. The arguments may only be of Python value types (
str) or of a subclass of
- A static method called
namethat returns the name of the branch (for use on the command line).
- A static method called
descriptionthat returns a description of the branch (for use on the command line).
Finally, register this branch in
4.10 Adding a New Platform
Platform class (from
platforms/base.py) in a new file in the
platforms directory. Be sure to make it a dataclass. Add any additional fields and, optionally, help strings for these fields (named
_f for a field
f). Implement all the required abstract properties (
imagenet_root if ImageNet is available). Finally, override
run_job() if different behavior is needed for the platform; be sure to ensure that the modified
run_job() method still installs the platform instance before calling the job function
Thank you to Ari Morcos and David Schwab for supporting the development of this framework and the research that we conducted with it. Thank you to FAIR for allowing me to open-source this framework. Thank you to David Bieber for teaching me how to do software engineering around deep learning and for the many ideas I borrowed from Python Fire.