Xbrain Unets Optimization

This example demonstrates how to run the optimization pipeline for the xbrain unets training job. It also describes how to add custom optimization algorithms to the conduit framework. Works in the optimization_gt branch.

Workflow description

This simple workflow trains a u-net classifier for segmenting X-brain data. Details on u-nets can be found here: https://arxiv.org/abs/1505.04597, and have been extensively used for classification. In this case, the network is trained to match ground truth annotations, and an f1-score is produced. This workflow searches over the available hyperparameters for this tool to find the best training accuracy.

Available Hyperparameters

• n_epochs: integer number of epochs for training
• mb_size: size of minibatches during training
• n_mb_per_epoch: number of minibatches per epoch
• tile_size: must be a multiple of 16, is the size of the tiles fed into the first layer of network (and determines network size)

If using SGD

• learning_rate: learning rate for gradient descent
• decay: decay param for gradient descent
• momentum: momentum param for gradient descent

If using ADAM

• learning_rate: learning rate for ADAM
• beta1, beta2: adaptation parameters for ADAM

Building tool containers

If you have ran the xbrain unets example, this step isn't necessary.

1. Build the xbrain unets training container.

   cd saber/saber/xbrain/unets
   docker build -t aplbrain/xbrain:unets .

Setting workflow and optimization parameters.

CWL, job, and parameters files can be found in saber/xbrain/jobs/unet_train_job/. There is two sets of files: one that uses SGD (stochastic gradient descent) and another one that uses the ADAM optimizer for training. In this example, we will use the SGD files.

1. Enter your BOSS token in the xbrain_unets_ex_job_sgd.yml with access to the appropriate channels.
2. Since unets training is GPU enabled, we will be running a local workflow with intermediate storage to the local file system. More information about storage and local execution can be found on the local storage wiki. For larger optimization tasks, xbrain_unets_train_sgd.cwl can be configured to use Amazon S3. S3 access requires an AWS credentials file with the appropriate permissions. Instructions for setting up AWS can be found here.
3. params_sgd.yml is configured to optimize only the number of epochs from a small to reduce time per run. Verify that the correct parameters and bounds are set for the run.

Compiling and launching the workflow

This step is like the other workflows, but it doesn't require you to run the workflow through the web interface.

1. Open a bash shell in the cwl_parser container

   docker exec -it saber_cwl_parser_1 /bin/bash
   

2. Navigate to the xbrain unets job folder and run the conduit build command.

   cd saber/xbrain/jobs/unet_train_job
   conduit build xbrain_unets_train_sgd.cwl xbrain_unets_ex_job_sgd.yml
   

3. Now start the optimization with conduit optimize

   conduit optimize xbrain_unets_train_sgd.cwl xbrain_unets_ex_job_sgd.yml params_sgd.yml --output optiout.csv --max-iterations 5 --sampling-strategy random
   

   This will run a workflow, collect the F1 scores of the xbrain pipeline on the image, and then resample the parameters set the in the params.yml file and run another workflow. Finally, it will save the results of all the runs, along with their parameters, to optiout.csv

   The optimization method used by default is random, but can be expanded by the user.

4. If an iteration fails, detailed logs can be accessed in the volumes/logs/ directory or through the airflow web app.

Adding new sampling strategies

Currently, we only have random sampling in our optimization framework. Sampling strategies can be added by writing a class that inherits from the abstract class Sampler defined in conduit/utils/parameterize.py. The sampler should have a method for updating the parameters for the next run and a sampling generator.

A sampler can either create its own parameterization grid (iterable list containing the parameter space) or use the existing parameterize function that exists within parameterize.py.

Once the sampler is defined in parameterize.py, the user must add it as a recognized flag in the file conduit/conduit.

1. Find the optimize api call defined in conduit/conduit.
2. Below the lines:

if args.sampling_strategy == 'random':
    sampler = RandomSampler(pm, job, args.max_iterations)

Add an elif statement that initializes the custom sampler class. For example:

elif args.sampling_strategy == 'my_sampler':
   sampler = MySampler(pm, job, args.max_iterations)

Any additional parameters needed for the samplers would also need to be added in the optimize_parser arguments in conduit/conduit.