Running and configuration

Working directory

I generally store every untracked file in an untracked “working directory” called workdir, at the root of the project. It should live on a data drive (ie one with high capacity). On ds-serv2, I make it a softlink to a subdir of /data.

Here is its structure:

workdir
├── datasets
│   ├── configs # contains dataset configurations. consider moving this to a tracked directory
│   ├── processed # contains processed datasets
│   └── raw # contains raw datasets
└── models # contains checkpoints and training results

all of these directories will be automatically created when the code is run (except for workdir/datasets/configs. There is no reference in the code to workdir/datasets/configs either, so you can just move it wherever)

Config file

The program will read some values from a global config file at the root of the project called config.yaml.

Here is what I use (for the rest I use the default)

wandb_key: [the api key for weights and biases]
wandb_project: [the project name for weights and biases]
processed_dataset_bucket: [the gcp bucket containing processed datasets]
checkpoint_bucket: [the gcp bucket containing model checkpoints and training run info]

The full definition of all fields for the config file is in nugets.pipeline.config.GlobalConf.

Controlling which GPU is being used

By default the cuda:0 device will be used, but another one can be selected. This is done by setting the CUDA_VISIBLE_DEVICES environment variable in the shell.

e.g.

$ export CUDA_VISIBLE_DEVICES=2
$ ./nugets.sh [...] # should use GPU 2

Training

Using config files

Note

This section is about model config files. For the global config file, see the previous section.

You can use the train_from_config subcommand to use a config file to specify training parameters from a YAML file instead. A good location to store those could be workdir/configs for ex.

.venv/bin/python -m nugets train_from_config workdir/model_config.yaml

where model_config.yaml has the following structure

task:
   type: TaskName
   dataset: DatasetName
   dataset_config:
      dataset_parameter1: value1
      dataset_parameter2: value2
   ood_dataset: DatasetName
   ood_config:
      dataset_parameter1: value1
      dataset_parameter2: value2
backbone:
   type: BackboneName
   hyperparameter1: value1
   hyperparameter2: value2
   [...]
batch_size: 16
learning_rate: 1.e-5
debug_mode: False

Using a bash script: nugets.sh

I added a nugets.sh script that essentially

1. sets stuff in the the environment

2. runs the code with automated debugging (meaning the ipdb debugger will automatically run if an exception is not caught)

Note

This is not fully optimal in a production environment:

• running with automated debugging may have a small performance overhead (not sure this matters when running on gpu though)

• running this way makes torch-lightning think it is in an “interactive environment” (such as a jupyter notebook) because I am using ipython there. This makes it disable some features such as multi-gpu training (although I don’t think we want to enable that)

$ ./nugets.sh  --help
usage: __main__.py [-h] [--config CONFIG] [-d] [-v] [--logfile LOGFILE]
                {train,train_from_config,wandb_agent} ...

NUGETS - NeUral GEomeTry Suite

positional arguments: {train,train_from_config,wandb_agent}
    train               Train a model
    train_from_config   train a model from a config file
    wandb_agent         run the agent for a wandb sweep

options:
  -h, --help            show this help message and exit
  --config CONFIG       global config path
  -d, --debug           Print lots of debugging statements
  -v, --verbose         Be verbose
  --logfile LOGFILE     Log file to write to (implies --verbose)

Using Weights and biases sweeps

See also: Cloud integration and automatisation.

This is the solution for automation: you can run a W&B sweep worker by using the following command:

./nugets.sh wandb_agent $SWEEP_ID

Using only the command line (not recommanded)

You can use the train subcommand to specify all model parameters throguh the command line. This is not recommended because it leads to very long commands, this is mostly a result of one of the authors having fun with argparse, but it can be useful for quick testing (especially for models with not too many hyperparameters or with sensible defaults)

There is an example is an example for that in the test_train_command.sh file. The command looks like this:

.venv/bin/python -m nugets train \
    --batch-size 12 --learning-rate 1.e-6 \
    --task "WassersteinDistanceTask" \
    --dataset workdir/datasets/configs/dummy_growing_circles.yaml \
    --backbone-type CoupledNetwork  \
    --backbone-latent-dimension 256 \
    --backbone-p 2 \
    --backbone-decoder-distance SinkhornLoss \
    --backbone-encoder-backbone Transformer \
        --backbone-encoder-n-heads 12 \
        --backbone-encoder-n-layers 6 \
        --backbone-encoder-d-model 768 \
        --backbone-encoder-feed-forward-hidden-dim 51

Note

Here the dataset is still specified through a yaml config.