Background¶

I am getting my PhD in astrophysics from UC Berkeley. My research focuses on observational cosmology, which boils down to analyzing images of the night sky. I help develop a pipeline that automatically detects stars and galaxies in multi-wavelength images and fits models to them. From the sample of galaxies we collect, it is possible to answer questions like, “how has the expansion rate of the Universe changed over time?”. The Sloan Digital Sky Survey (SDSS) showed that you can do this with a large enough sample (about a million galaxies) and region of the sky (1/3 of it). My team is carrying out the Legacy Survey to detect 30x more, and up to 6x fainter, galaxies than SDSS. The most distant galaxies are about 10 billion lightyears away.

Our goal is to create a 2D map of the positions of about 30 million galaxies extracted from more than 100 TBs of images. Given the locations, we can take a spectrum of each galaxy (e.g. how bright it is at many wavelengths of visible light) and infer how far away it is. From this 3D map we can measure the expansion rate of the Universe at different points in time (Eisenstein & Hu 1998; Eisenstein et al. 2005; Seo & Eisenstein 2007; Butler et al. 2017).

The primary goal of my thesis is to measure the statistical bias and variance of our pipeline. How does our completeness depend on whether a galaxy is bright or faint, blue or red, big or small? How well does our pipeline handle image artifacts, instrument issues, or transient objects?

To answer all of this, John Moustakas, a professor at Siena College, and I developed the obiwan code. It does Monte Carlo simulations of the pipeline by adding fake galaxies to random locations in the images, running the pipeline, and repeating.

The success of all this depends on whether or not the fake galaxies are representative of the true galaxy population. If not, we are measuring the bias and variance of a sample that does not exist in the data. I need to show that the fake galaxies “look” like the real ones. Enter the Convolutional Neural Network (CNN)…

The Problem¶

This is a supervised binary classification problem of labeled multi-wavelength images. The training set is huge, at least 60 million examples, and each example is six images (64 x 64 x 6 pixels). There is one image for each of the three wavelength bands and another three images to encode camera artifacts and the variance of each pixel.

The CNN predicts whether a galaxy is fake (1) or real (0).

Novel¶

This project is novel because half the training images are model-generated and the goal is for the CNN to do poorly. If the CNN cannot do better than guessing the most numerous class, then the fake galaxies “look like” real ones. i.e. my model is representative of reality. This has generated interest in the team since a similar CNN could be trained using models of other types of galaxies.

Examples of Training Data¶

The galaxies below are some of the brightest in the training set (99th percentile in brightness). More Examples of fainter galaxies (75th, 50th, 25th, and 1st percentiles in brightness) are at the end.

Figure 1. The label for each image is on the left (R for Real and F for Fake) and its corresponding g-band magnitude is the number on the right (the smaller the number, the brighter the galaxy). Each row represents a single galaxy imaged at three different wavelengths. The color-image (left most panel) shows the colors you would see by eye, while the black and white-images (right six panels) are the training data: three images for the g, r, and z-bands and their corresponding artifact/inverse variance (ivar) images. Consecutive rows of R and F (rows 1 and 2, 3 and 4, etc.) have similar g-band magnitudes so that a fair comparison can be made.

The examples reveal at least two challenges for the CNN.

1. Only the central-object matters, but there are many off-center objects in the images. These are random background sources, often bright galaxies or stars that we are not interested in.
2. These galaxies are very faint. The CNN must be able to dig out the low Signal to Noise sources.

CNN Architecture¶

As a starting point, I used TensorFlow to build a CNN similar to LeNet-5 with the following architecture:

The input image has 64 x 64 x 6 pixels. The CNN is much shallower than the ImageNet ILSVRC winners, so in addition to tuning the number of feature maps, kernel size, stride, etc., I plan to make it deeper.

Intel Xeon Phi¶

I have used the Cray XC30 (Edison) and Cray XC40 (Cori) supercomputers at the National Energy Research Scientific Computing Center (NERSC) for the majority of my thesis work. With almost 10,000 Intel Xeon Phi processor nodes on Cori, NERSC Staff are particularly interested in helping users run on Xeon Phi.

I decided to train on Xeon Phi, instead of a GPU, when NERSC/Intel released optimized installs of many of the popular machine learning libraries (Caffe, TensorFlow, Theano, Torch, see full list). I created an initial training set of 2048 images with an equal number of fake and real galaxies. The images are float32 so I stored every 512 examples in a file, thinking that a 50 MB file would fit in memory on most machines.

It takes about 3 minutes to train 4 epochs of 2048 images on a single Xeon Phi node (68 hardware cores). For hundreds of nodes, I plan on training a different minibatch with each MPI task, updating a global set of weights using the learned weights from the minibatches, then repeating. Although multi-node support is not yet available to users, NERSC Staff can scale ResNet-50 and DCGAN to 1024 Xeon Phi nodes.

Fortunately, I hope to begin multi-node training soon as the NERSC Staff have volunteered my CNN for non-benchmark multi-node testing.

TensorBoard & Profiling¶

The accuracy, loss, and graph from the 4 training epochs are shown below using TensorBoard. Different colors correspond to me restarting the training twice to demonstrate that the checkpoints are working.

Figure 2. Accuracy and loss with TensorBoard

Figure 3. Graph with TensorBoard

I also profile my CNN using TensorFlow’s timeline object. This writes timings for each node of the graph to a json file. You can inspect it with Google Chrome, by going to chrome://tracing and clicking load. Here’s what it looks like.

Figure 4. Profiling with TensorFlow’s timeline and Google Chrome

Instructions¶

See the following for training the CNN at NERSC.

#SBATCH -p regular
#SBATCH -N 50
#SBATCH -t 01:00:00
#SBATCH --account=desi
#SBATCH -J train
#SBATCH -L SCRATCH,project
#SBATCH -C haswell

let tasks=32*${SLURM_JOB_NUM_NODES}

# NERSC / Cray / Cori / Cori KNL things
export KMP_AFFINITY=disabled
export MPICH_GNI_FORK_MODE=FULLCOPY
export MKL_NUM_THREADS=1
export OMP_NUM_THREADS=1

srun -n ${tasks} -c 1 python create_training.py \
     --which tractor --bricks_fn bricks.txt --nproc ${tasks}
     --savedir /global/cscratch1/sd/kaylanb/obiwan_out/dr5_hdf5

srun -n ${tasks} -c 1 python split_testtrain.py \
     --bricks_fn bricks.txt --nproc ${tasks} \
     --real_dir /global/cscratch1/sd/kaylanb/obiwan_out/dr5_hdf5 \
     --sim_dir /global/cscratch1/sd/kaylanb/obiwan_out/elg_dr5_coadds \
     --save_dir /global/cscratch1/sd/kaylanb/obiwan_out/dr5_testtrain

More Examples¶

Galaxies with 75th perentile in brightness:

Galaxies with 50th perentile in brightness:

Galaxies with 25th perentile in brightness:

Galaxies with 1st percentile in brightness (faintest in the training set):