L46 Datamaps

This is the codebase for the paper Knowledge Distillation with Training Dynamics, Hromadka & Rex.

In this study, we examine the efficacy of Knowledge Distillation in conjunction with Dataset Cartography within Computer Vision. Our aim was to ascertain whether training on subsets of data, selected based on training dynamics, can enhance the performance of student models while being cost-effective. Contrary to expectations set by similar NLP studies, our experiments across various architectures and datasets (CIFAR10, CIFAR100, MNIST) revealed that randomly selected data subsets (approximately 33% of the total dataset) consistently outperformed those chosen based on specific training dynamics. This unexpected outcome suggests the need for a nuanced understanding of dataset selection in Knowledge Distillation, particularly in the Computer Vision domain. Our findings highlight the complex interplay between data selection and model performance, opening new avenues for research in effective training strategies. We publish our code for future research.

Installing the project

You must have conda installed locally. All project dependencies are included in the conda file: environment.yml.

This was tested to work on: CUDA Version: 11.2 (NVIDIA-SMI 460.32.03).

<!-- Install the codebase -->
cd REPOSITORY
conda create python=3.10.12 --name L46
conda activate L46

<!-- Install PyTorch if not already on the system -->
conda install pytorch -c pytorch

<!-- Install requirements -->
pip install -r requirements.txt

<!-- Optionally, install lightgbm -->
# pip install lightgbm --install-option=--gpu --install-option="--opencl-include-dir=/usr/local/cuda/include/" --install-option="--opencl-library=/usr/local/cuda/lib64/libOpenCL.so"

Change BASE_DIR from /src/_config.py to point to the project directory on your machine.

Search am2770in the codebase and replace the paths with the paths on your machine.

Test your install by running:

./run_exp_test.sh

When you install new modules, save the environment:

pip freeze > requirements.txt

Codebase basics

• Store training metrics using wandb. It makes running experiments much easier.
• Train using pytorch-lightning, which is a wrapper on top of pytorch. It has three main benefits: (1) it make training logic much simpler, (2) it allows us to use the same code for training on CPU or GPU, and (3) integrates nicely with wandb.

Code structure

• src

  • run_experiment.py: code for parsing arguments, and starting experiment

    • def parse_arguments - include all command line arguments
    • def run_experiment - create wandb logger and start training model
    • def train_model - training neural networks using pytorch lightning
    • important commandline arguments
      • dataset
      • model
      • feature_extractor_dims - the size of the hidden layers in the dnn
      • max_steps - maximum training iterations
      • batchnorm, dropout_rate
      • lr, batch_size, patience_early_stopping
      • lr_scheduler - learning rate scheduler (the default scheduler reduces the learning rate until it gets 10x smaller)

  • dataset.py: loading the datasets

  • model.py: neural network architecture

    • def create_model - logic to create a neural network
    • class TrainingLightningModule - generic class that includes logging losses etc using pytorch-lightning
    • class DNN - neural network architecture. It's a very extensible class, that includes more things that you'll need.
    • class FirstLinearLayer - used for adding different first layers (e.g., DietNetworks, ConcreteLayer etc)
    • class WeightPredictorNetwork - for DietNetworks
• scripts_experiments
  • template to set a wandb sweep
• compute_results
  • template to retrieve results from wandb (using SQL-style notation), and analyse them easily
• some utilities
  • run_exp_test.sh - run a test experiment
  • generate_uuid.py - generate a unique id for an experiment
  • kill_tmux_session.sh $1 - kill a specific tmux session

Changing to a specific branch

git pull origin BRANCH
git checkout BRANCH

WandB sweep

• note: the wandb sweep run on the current codebase. Thus, if you change the code while running the sweep, new runs will run on the new codebase.

wandb sweep scripts_experiments/_template_sweep.yaml
wandb agent XYZ

Reproducing Experiments

WandB sweep yaml files can be found in /scripts_experiments. These correspond to experimental phases of the project, and can be ran with wandb sweep in order to repeat our experiments.