Repository: neuronets/nobrainer

Issue:
Using nobrainer model for binary classification
Hi, 

I am trying to use the nobrainer pre-trained model for a binary classification problem instead of 3D segmentation. Essentially I want my model to learn brain information (structure, boundaries etc) from the pre-trained nobrainer model and then use that to further train a model for classification. I am using the ABIDE dataset and currently have binary labels for each volume. From what I understood after taking a look at the `convert()` method in io.py, it takes in an input of filepaths in pairs of 2 and both of them need to be brain volumes. Is there a way for me to bypass giving the labels as 3D volumes and still use the brain volume conformation just on the features (ABIDE volumes). 

As a little background, I am clipping the model at the the end of the encoder path and attaching my own layers. I want to conform the variable shaped dataset to fixed shapes of (128, 128, 128) blocks and pass them through the nobrainer model (with the weights set from your pre-trained model) and then visualize the activations layers to understand what its learning and if the model is viable option for my goal.

Discussion:

Patch:
diff --git a/.gitignore b/.gitignore
index 8ce86f7..0d10589 100644
--- a/.gitignore
+++ b/.gitignore
@@ -10,4 +10,6 @@ __pycache__
 *.pytest_cache
 *.sif
 *.simg
+*.tfrec
 *.tfrecords
+scratch
diff --git a/.travis.yml b/.travis.yml
index 0e5bf08..73674c1 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -8,4 +8,4 @@ branches:
 install:
   - pip install --no-cache-dir --editable=".[cpu,dev]"
 script:
-  - pytest --verbose --cov=nobrainer nobrainer
+  - travis_wait pytest --verbose --cov=nobrainer nobrainer
diff --git a/guide/01-getting_started.ipynb b/guide/01-getting_started.ipynb
index 11953ae..76b4aa4 100644
--- a/guide/01-getting_started.ipynb
+++ b/guide/01-getting_started.ipynb
@@ -99,12 +99,12 @@
     "### Layout\n",
     "\n",
     "- `nobrainer.io`: input/output methods\n",
-    "- `nobrainer.layers`: custom layers, which conform to the Keras API\n",
+    "- `nobrainer.layers`: custom Keras layers\n",
     "- `nobrainer.losses`: loss functions for volumetric segmentation\n",
     "- `nobrainer.metrics`: metrics for volumetric segmentation\n",
     "- `nobrainer.models`: pre-defined Keras models\n",
-    "- `nobrainer.training`: training utilities (supports training on single and multiple GPUs)\n",
-    "- `nobrainer.transform`: random rigid transformations for data augmentation\n",
+    "- `nobrainer.tfrecords`: writing and reading of TFRecords files\n",
+    "- `nobrainer.transform`: rigid transformations for data augmentation\n",
     "- `nobrainer.volume`: `tf.data.Dataset` creation and data augmentation utilities"
    ]
   },
@@ -136,7 +136,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.7"
+   "version": "3.7.6"
   }
  },
  "nbformat": 4,
diff --git a/guide/02-preparing_training_data.ipynb b/guide/02-preparing_training_data.ipynb
index abb2edb..b4b002b 100644
--- a/guide/02-preparing_training_data.ipynb
+++ b/guide/02-preparing_training_data.ipynb
@@ -96,10 +96,9 @@
    "source": [
     "!nobrainer convert \\\n",
     "    --csv='/tmp/nobrainer-data/filepaths.csv' \\\n",
-    "    --tfrecords-template='tfrecords/data_shard-{shard:03d}.tfrecords' \\\n",
-    "    --volumes-per-shard=4 \\\n",
+    "    --tfrecords-template='tfrecords/data_shard-{shard:03d}.tfrec' \\\n",
+    "    --examples-per-shard=2 \\\n",
     "    --volume-shape 256 256 256 \\\n",
-    "    --num-parallel-calls=2 \\\n",
     "    --verbose"
    ]
   },
@@ -129,7 +128,7 @@
    "outputs": [],
    "source": [
     "# A glob pattern to match the files we want to train on.\n",
-    "file_pattern = 'tfrecords/data_shard-*.tfrecords'\n",
+    "file_pattern = 'tfrecords/data_shard-*.tfrec'\n",
     "\n",
     "# The number of classes the model predicts. A value of 1 means the model performs\n",
     "# binary classification (i.e., target vs background).\n",
@@ -192,6 +191,13 @@
     "dataset"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -308,7 +314,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.7"
+   "version": "3.7.6"
   }
  },
  "nbformat": 4,
diff --git a/guide/train_binary_classification.ipynb b/guide/train_binary_classification.ipynb
new file mode 100644
index 0000000..dac8179
--- /dev/null
+++ b/guide/train_binary_classification.ipynb
@@ -0,0 +1,389 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Train a neural network for binary classification\n",
+    "\n",
+    "In this notebook, we will use Nobrainer to train a model to classify brains into one of two classes. This task is different from segmentation -- the model will output a single scalar per brain.\n",
+    "\n",
+    "In the following cells, we will:\n",
+    "\n",
+    "1. Get sample T1-weighted MR scans as features and random boolean values as labels.\n",
+    "2. Convert the data to TFRecords format.\n",
+    "3. Create two Datasets of the features and labels.\n",
+    "    - One dataset will be for training and the other will be for evaluation.\n",
+    "4. Instantiate a 3D convolutional neural network.\n",
+    "5. Choose a loss function and metrics to use.\n",
+    "6. Train on part of the data.\n",
+    "7. Evaluate on the rest of the data.\n",
+    "\n",
+    "## Google Colaboratory\n",
+    "\n",
+    "If you are using Colab, please switch your runtime to GPU. To do this, select `Runtime > Change runtime type` in the top menu. Then select GPU under `Hardware accelerator`. A GPU greatly speeds up training."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import nobrainer"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Get sample features and labels\n",
+    "\n",
+    "We use 9 pairs of volumes for training and 1 pair of volumes for evaulation. Many more volumes would be required to train a model for any useful purpose."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import random\n",
+    "\n",
+    "csv_of_filepaths = nobrainer.utils.get_data()\n",
+    "filepaths = nobrainer.io.read_csv(csv_of_filepaths)\n",
+    "\n",
+    "# Add random boolean values\n",
+    "filepaths = [(x, random.choice([0, 1])) for x, _ in filepaths]\n",
+    "\n",
+    "train_paths = filepaths[:9]\n",
+    "evaluate_paths = filepaths[9:]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "filepaths"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Convert medical images to TFRecords\n",
+    "\n",
+    "Remember how many full volumes are in the TFRecords files. This will be necessary to know how many steps are in on training epoch. The default training method needs to know this number, because Datasets don't always know how many items they contain."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Verify that all volumes have the same shape and that labels are integer-ish.\n",
+    "\n",
+    "invalid = nobrainer.io.verify_features_labels(train_paths)\n",
+    "assert not invalid\n",
+    "\n",
+    "invalid = nobrainer.io.verify_features_labels(evaluate_paths)\n",
+    "assert not invalid"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!mkdir -p data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "# Convert training and evaluation data to TFRecords.\n",
+    "\n",
+    "nobrainer.tfrecord.write(\n",
+    "    features_labels=train_paths,\n",
+    "    filename_template='data/data-train_shard-{shard:03d}.tfrec',\n",
+    "    examples_per_shard=3)\n",
+    "\n",
+    "nobrainer.tfrecord.write(\n",
+    "    features_labels=evaluate_paths,\n",
+    "    filename_template='data/data-evaluate_shard-{shard:03d}.tfrec',\n",
+    "    examples_per_shard=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!ls data"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Create Datasets"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "n_classes = 1\n",
+    "batch_size = 2\n",
+    "volume_shape = (256, 256, 256)\n",
+    "block_shape = (128, 128, 128)\n",
+    "n_epochs = None\n",
+    "num_parallel_calls = 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import tensorflow as tf\n",
+    "\n",
+    "def get_dataset(file_pattern, shuffle=True):\n",
+    "    d = tf.data.Dataset.list_files(file_pattern, shuffle=shuffle)\n",
+    "    d = d.interleave(\n",
+    "        map_func=lambda x: tf.data.TFRecordDataset(x, compression_type=\"GZIP\"),\n",
+    "        cycle_length=num_parallel_calls,\n",
+    "        num_parallel_calls=num_parallel_calls)\n",
+    "\n",
+    "    # Parse each example in each TFRecords file as a tensor of features and a\n",
+    "    # tensor of labels.\n",
+    "    parse_fn = nobrainer.tfrecord.parse_example_fn(\n",
+    "        volume_shape=volume_shape, scalar_label=True)\n",
+    "    d = d.map(map_func=parse_fn)\n",
+    "    return d\n",
+    "\n",
+    "def _preprocess(x, y):\n",
+    "    x = nobrainer.volume.standardize(x)\n",
+    "    x = nobrainer.volume.to_blocks(x, block_shape=block_shape)\n",
+    "    n_blocks = tf.shape(x)[0]\n",
+    "    # Add grayscale channel.\n",
+    "    return tf.expand_dims(x, axis=-1), tf.repeat(y, n_blocks)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset_train = get_dataset(\"data/data-train*.tfrec\")\n",
+    "dataset_train = dataset_train.map(_preprocess)\n",
+    "# We need to unbatch because separating into blocks added a dimension.\n",
+    "dataset_train = dataset_train.unbatch()\n",
+    "dataset_train = dataset_train.batch(batch_size, drop_remainder=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset_evaluate = get_dataset(\"data/data-evaluate*.tfrec\", shuffle=False)\n",
+    "dataset_evaluate = dataset_evaluate.map(_preprocess)\n",
+    "# We need to unbatch because separating into blocks added a dimension.\n",
+    "dataset_evaluate = dataset_evaluate.unbatch()\n",
+    "dataset_evaluate = dataset_evaluate.batch(batch_size, drop_remainder=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset_train"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset_evaluate"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Instantiate a neural network"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = tf.keras.Sequential([\n",
+    "    tf.keras.layers.Conv3D(filters=12, kernel_size=3, padding=\"same\", \n",
+    "                           strides=2, activation=\"relu\", input_shape=(*block_shape, 1)),\n",
+    "    tf.keras.layers.Flatten(),\n",
+    "    tf.keras.layers.Dense(128),\n",
+    "    tf.keras.layers.Dense(n_classes, activation=\"sigmoid\")\n",
+    "])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "model.summary()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Choose a loss function and metrics"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import tensorflow as tf"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "optimizer = tf.keras.optimizers.Adam(learning_rate=1e-04)\n",
+    "\n",
+    "model.compile(\n",
+    "    optimizer=optimizer,\n",
+    "    loss=tf.keras.losses.binary_crossentropy,\n",
+    "    metrics=[\"accuracy\"],\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Train and evaluate model\n",
+    "\n",
+    "$$\n",
+    "steps = \\frac{nBlocks}{volume} * \\frac{nVolumes}{batchSize}\n",
+    "$$"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "steps_per_epoch = nobrainer.volume.get_steps_per_epoch(\n",
+    "    n_volumes=len(train_paths),\n",
+    "    volume_shape=volume_shape,\n",
+    "    block_shape=block_shape,\n",
+    "    batch_size=batch_size)\n",
+    "\n",
+    "steps_per_epoch"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "validation_steps = nobrainer.volume.get_steps_per_epoch(\n",
+    "    n_volumes=len(evaluate_paths),\n",
+    "    volume_shape=volume_shape,\n",
+    "    block_shape=block_shape,\n",
+    "    batch_size=batch_size)\n",
+    "\n",
+    "validation_steps"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "model.fit(\n",
+    "    dataset_train,\n",
+    "    epochs=5,\n",
+    "    steps_per_epoch=steps_per_epoch, \n",
+    "    validation_data=dataset_evaluate, \n",
+    "    validation_steps=validation_steps)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Predict natively medical images (without TFRecords)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# TODO"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python [conda env:nb]",
+   "language": "python",
+   "name": "conda-env-nb-py"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/guide/train_binary_segmentation.ipynb b/guide/train_binary_segmentation.ipynb
index 5651486..f71a4b3 100644
--- a/guide/train_binary_segmentation.ipynb
+++ b/guide/train_binary_segmentation.ipynb
@@ -25,22 +25,6 @@
     "If you are using Colab, please switch your runtime to GPU. To do this, select `Runtime > Change runtime type` in the top menu. Then select GPU under `Hardware accelerator`. A GPU greatly speeds up training."
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Install nobrainer"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "!pip install --no-cache-dir nobrainer[gpu]"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -89,7 +73,7 @@
    "source": [
     "# Verify that all volumes have the same shape and that labels are integer-ish.\n",
     "\n",
-    "invalid = nobrainer.io.verify_features_labels(train_paths, num_parallel_calls=2)\n",
+    "invalid = nobrainer.io.verify_features_labels(train_paths)\n",
     "assert not invalid\n",
     "\n",
     "invalid = nobrainer.io.verify_features_labels(evaluate_paths)\n",
@@ -108,22 +92,22 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "scrolled": true
+   },
    "outputs": [],
    "source": [
     "# Convert training and evaluation data to TFRecords.\n",
     "\n",
-    "nobrainer.io.convert(\n",
-    "    train_paths,\n",
-    "    tfrecords_template='data/data-train_shard-{shard:03d}.tfrecords',\n",
-    "    volumes_per_shard=3,\n",
-    "    num_parallel_calls=2)\n",
+    "nobrainer.tfrecord.write(\n",
+    "    features_labels=train_paths,\n",
+    "    filename_template='data/data-train_shard-{shard:03d}.tfrec',\n",
+    "    examples_per_shard=3)\n",
     "\n",
-    "nobrainer.io.convert(\n",
-    "    evaluate_paths,\n",
-    "    tfrecords_template='data/data-evaluate_shard-{shard:03d}.tfrecords',\n",
-    "    num_parallel_calls=1,\n",
-    ")"
+    "nobrainer.tfrecord.write(\n",
+    "    features_labels=evaluate_paths,\n",
+    "    filename_template='data/data-evaluate_shard-{shard:03d}.tfrec',\n",
+    "    examples_per_shard=1)"
    ]
   },
   {
@@ -165,7 +149,7 @@
    "outputs": [],
    "source": [
     "dataset_train = nobrainer.volume.get_dataset(\n",
-    "    file_pattern='data/data-train_shard-*.tfrecords',\n",
+    "    file_pattern='data/data-train_shard-*.tfrec',\n",
     "    n_classes=n_classes,\n",
     "    batch_size=batch_size,\n",
     "    volume_shape=volume_shape,\n",
@@ -177,7 +161,7 @@
     ")\n",
     "\n",
     "dataset_evaluate = nobrainer.volume.get_dataset(\n",
-    "    file_pattern='data/data-evaluate_shard-*.tfrecords',\n",
+    "    file_pattern='data/data-evaluate_shard-*.tfrec',\n",
     "    n_classes=n_classes,\n",
     "    batch_size=batch_size,\n",
     "    volume_shape=volume_shape,\n",
@@ -369,9 +353,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python [conda env:nb]",
    "language": "python",
-   "name": "python3"
+   "name": "conda-env-nb-py"
   },
   "language_info": {
    "codemirror_mode": {
@@ -383,7 +367,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.5.2"
+   "version": "3.7.6"
   }
  },
  "nbformat": 4,
diff --git a/guide/train_on_multiple_gpus.ipynb b/guide/train_on_multiple_gpus.ipynb
index 67e6e33..b759dc9 100644
--- a/guide/train_on_multiple_gpus.ipynb
+++ b/guide/train_on_multiple_gpus.ipynb
@@ -25,22 +25,6 @@
     "If you are using Colab, please switch your runtime to GPU. To do this, select `Runtime > Change runtime type` in the top menu. Then select GPU under `Hardware accelerator`. A GPU is not necessary to prepare the data, but a GPU is helpful for training a model, which we demonstrate at the end of this notebook. This will give you access to one GPU, but the code will still run properly. To actually train a model on multiple GPUs, you will have to use Cloud services, a high-performance computing cluster, or your own hardware."
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Install nobrainer"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "!pip install --no-cache-dir nobrainer[gpu]"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -113,17 +97,15 @@
    "source": [
     "# Convert training and evaluation data to TFRecords.\n",
     "\n",
-    "nobrainer.io.convert(\n",
-    "    train_paths,\n",
-    "    tfrecords_template='data/data-train_shard-{shard:03d}.tfrecords',\n",
-    "    volumes_per_shard=3,\n",
-    "    num_parallel_calls=2)\n",
+    "nobrainer.tfrecord.write(\n",
+    "    features_labels=train_paths,\n",
+    "    filename_template='data/data-train_shard-{shard:03d}.tfrec',\n",
+    "    examples_per_shard=3)\n",
     "\n",
-    "nobrainer.io.convert(\n",
-    "    evaluate_paths,\n",
-    "    tfrecords_template='data/data-evaluate_shard-{shard:03d}.tfrecords',\n",
-    "    num_parallel_calls=1,\n",
-    ")"
+    "nobrainer.tfrecord.write(\n",
+    "    features_labels=evaluate_paths,\n",
+    "    filename_template='data/data-evaluate_shard-{shard:03d}.tfrec',\n",
+    "    examples_per_shard=1)"
    ]
   },
   {
@@ -167,7 +149,7 @@
    "outputs": [],
    "source": [
     "dataset_train = nobrainer.volume.get_dataset(\n",
-    "    file_pattern='data/data-train_shard-*.tfrecords',\n",
+    "    file_pattern='data/data-train_shard-*.tfrec',\n",
     "    n_classes=n_classes,\n",
     "    batch_size=batch_size,\n",
     "    volume_shape=volume_shape,\n",
@@ -179,7 +161,7 @@
     ")\n",
     "\n",
     "dataset_evaluate = nobrainer.volume.get_dataset(\n",
-    "    file_pattern='data/data-evaluate_shard-*.tfrecords',\n",
+    "    file_pattern='data/data-evaluate_shard-*.tfrec',\n",
     "    n_classes=n_classes,\n",
     "    batch_size=batch_size,\n",
     "    volume_shape=volume_shape,\n",
@@ -304,9 +286,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python [conda env:nb]",
    "language": "python",
-   "name": "python3"
+   "name": "conda-env-nb-py"
   },
   "language_info": {
    "codemirror_mode": {
@@ -318,7 +300,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.5.2"
+   "version": "3.7.6"
   }
  },
  "nbformat": 4,
diff --git a/guide/train_on_tpu.ipynb b/guide/train_on_tpu.ipynb
index 70f1009..7612ab1 100644
--- a/guide/train_on_tpu.ipynb
+++ b/guide/train_on_tpu.ipynb
@@ -25,24 +25,6 @@
     "If you are using Colab, please switch your runtime to TPU. To do this, select `Runtime > Change runtime type` in the top menu. Then select TPU under `Hardware accelerator`."
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Install nobrainer and TensorFlow 1.13.1\n",
-    "\n",
-    "At the time of writing (April 17, 2019), running a Keras model on TPU works reliably with TensorFlow 1.13.1 but not with the nightly releases."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "!pip install --no-cache-dir nobrainer 'tensorflow-gpu==1.13.1'"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -115,17 +97,15 @@
    "source": [
     "# Convert training and evaluation data to TFRecords.\n",
     "\n",
-    "nobrainer.io.convert(\n",
-    "    train_paths,\n",
-    "    tfrecords_template='data/data-train_shard-{shard:03d}.tfrecords',\n",
-    "    volumes_per_shard=3,\n",
-    "    num_parallel_calls=2)\n",
+    "nobrainer.tfrecord.write(\n",
+    "    features_labels=train_paths,\n",
+    "    filename_template='data/data-train_shard-{shard:03d}.tfrec',\n",
+    "    examples_per_shard=3)\n",
     "\n",
-    "nobrainer.io.convert(\n",
-    "    evaluate_paths,\n",
-    "    tfrecords_template='data/data-evaluate_shard-{shard:03d}.tfrecords',\n",
-    "    num_parallel_calls=1,\n",
-    ")"
+    "nobrainer.tfrecord.write(\n",
+    "    features_labels=evaluate_paths,\n",
+    "    filename_template='data/data-evaluate_shard-{shard:03d}.tfrec',\n",
+    "    examples_per_shard=1)"
    ]
   },
   {
@@ -169,7 +149,7 @@
    "outputs": [],
    "source": [
     "dataset_train = nobrainer.volume.get_dataset(\n",
-    "    file_pattern='data/data-train_shard-*.tfrecords',\n",
+    "    file_pattern='data/data-train_shard-*.tfrec',\n",
     "    n_classes=n_classes,\n",
     "    batch_size=batch_size,\n",
     "    volume_shape=volume_shape,\n",
@@ -181,7 +161,7 @@
     ")\n",
     "\n",
     "dataset_evaluate = nobrainer.volume.get_dataset(\n",
-    "    file_pattern='data/data-evaluate_shard-*.tfrecords',\n",
+    "    file_pattern='data/data-evaluate_shard-*.tfrec',\n",
     "    n_classes=n_classes,\n",
     "    batch_size=batch_size,\n",
     "    volume_shape=volume_shape,\n",
@@ -316,9 +296,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python [conda env:nb]",
    "language": "python",
-   "name": "python3"
+   "name": "conda-env-nb-py"
   },
   "language_info": {
    "codemirror_mode": {
@@ -330,7 +310,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.5.2"
+   "version": "3.7.6"
   }
  },
  "nbformat": 4,
diff --git a/guide/transfer_learning.ipynb b/guide/transfer_learning.ipynb
index 60b21d2..f2658c8 100644
--- a/guide/transfer_learning.ipynb
+++ b/guide/transfer_learning.ipynb
@@ -27,22 +27,6 @@
     "If you are using Colab, please switch your runtime to GPU. To do this, select `Runtime > Change runtime type` in the top menu. Then select GPU under `Hardware accelerator`. A GPU greatly speeds up training."
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Install nobrainer"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "!pip install --no-cache-dir nobrainer[gpu]"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -115,17 +99,15 @@
    "source": [
     "# Convert training and evaluation data to TFRecords.\n",
     "\n",
-    "nobrainer.io.convert(\n",
-    "    train_paths,\n",
-    "    tfrecords_template='data/data-train_shard-{shard:03d}.tfrecords',\n",
-    "    volumes_per_shard=3,\n",
-    "    num_parallel_calls=2)\n",
+    "nobrainer.tfrecord.write(\n",
+    "    features_labels=train_paths,\n",
+    "    filename_template='data/data-train_shard-{shard:03d}.tfrec',\n",
+    "    examples_per_shard=3)\n",
     "\n",
-    "nobrainer.io.convert(\n",
-    "    evaluate_paths,\n",
-    "    tfrecords_template='data/data-evaluate_shard-{shard:03d}.tfrecords',\n",
-    "    num_parallel_calls=1,\n",
-    ")"
+    "nobrainer.tfrecord.write(\n",
+    "    features_labels=evaluate_paths,\n",
+    "    filename_template='data/data-evaluate_shard-{shard:03d}.tfrec',\n",
+    "    examples_per_shard=1)"
    ]
   },
   {
@@ -167,7 +149,7 @@
    "outputs": [],
    "source": [
     "dataset_train = nobrainer.volume.get_dataset(\n",
-    "    file_pattern='data/data-train_shard-*.tfrecords',\n",
+    "    file_pattern='data/data-train_shard-*.tfrec',\n",
     "    n_classes=n_classes,\n",
     "    batch_size=batch_size,\n",
     "    volume_shape=volume_shape,\n",
@@ -179,7 +161,7 @@
     ")\n",
     "\n",
     "dataset_evaluate = nobrainer.volume.get_dataset(\n",
-    "    file_pattern='data/data-evaluate_shard-*.tfrecords',\n",
+    "    file_pattern='data/data-evaluate_shard-*.tfrec',\n",
     "    n_classes=n_classes,\n",
     "    batch_size=batch_size,\n",
     "    volume_shape=volume_shape,\n",
@@ -359,9 +341,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python [conda env:nb]",
    "language": "python",
-   "name": "python3"
+   "name": "conda-env-nb-py"
   },
   "language_info": {
    "codemirror_mode": {
@@ -373,7 +355,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.3"
+   "version": "3.7.6"
   }
  },
  "nbformat": 4,
diff --git a/guide/variational_weights.ipynb b/guide/variational_weights.ipynb
index 0d84be4..ef62fe9 100644
--- a/guide/variational_weights.ipynb
+++ b/guide/variational_weights.ipynb
@@ -44,7 +44,7 @@
     "block_shape = (64, 64, 64)\n",
     "\n",
     "dataset = nobrainer.volume.get_dataset(\n",
-    "    file_pattern='tfrecords/data_shard-*.tfrecords',\n",
+    "    file_pattern='tfrecords/data_shard-*.tfrec',\n",
     "    n_classes=n_classes,\n",
     "    batch_size=batch_size,\n",
     "    volume_shape=volume_shape,\n",
@@ -161,9 +161,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python [conda env:nb]",
    "language": "python",
-   "name": "python3"
+   "name": "conda-env-nb-py"
   },
   "language_info": {
    "codemirror_mode": {
@@ -175,7 +175,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.5.2"
+   "version": "3.7.6"
   }
  },
  "nbformat": 4,
diff --git a/nobrainer/__init__.py b/nobrainer/__init__.py
index 64e067a..7227b83 100644
--- a/nobrainer/__init__.py
+++ b/nobrainer/__init__.py
@@ -18,7 +18,6 @@ import nobrainer.losses
 import nobrainer.metrics
 import nobrainer.models
 import nobrainer.prediction
-import nobrainer.training
 import nobrainer.transform
 import nobrainer.utils
 import nobrainer.volume
diff --git a/nobrainer/cli/main.py b/nobrainer/cli/main.py
index e81b651..a755ef8 100644
--- a/nobrainer/cli/main.py
+++ b/nobrainer/cli/main.py
@@ -15,14 +15,12 @@ import skimage.transform
 import tensorflow as tf
 
 from nobrainer import __version__
-from nobrainer.io import convert as _convert
 from nobrainer.io import verify_features_labels as _verify_features_labels
 from nobrainer.io import read_csv as _read_csv
 from nobrainer.io import read_volume as _read_volume
 from nobrainer.losses import get as _get_loss
 from nobrainer.prediction import _transform_and_predict
-from nobrainer.training import train as _train
-from nobrainer.utils import _get_all_cpus
+from nobrainer.tfrecord import write as _write_tfrecord
 from nobrainer.volume import from_blocks_numpy as _from_blocks_numpy
 from nobrainer.volume import get_dataset as _get_dataset
 from nobrainer.volume import get_steps_per_epoch as _get_steps_per_epoch
@@ -56,22 +54,26 @@ def cli():
 @click.option('-c', '--csv', type=click.Path(exists=True), required=True, **_option_kwds)
 @click.option('-t', '--tfrecords-template', default='tfrecords/data_shard-{shard:03d}.tfrecords', required=True, **_option_kwds)
 @click.option('-s', '--volume-shape', nargs=3, type=int, required=True, **_option_kwds)
-@click.option('-n', '--volumes-per-shard', type=int, default=100, help='Number of volume pairs per TFRecords file.', **_option_kwds)
+@click.option('-n', '--examples-per-shard', type=int, default=100, help='Number of (feature, label) pairs per TFRecord file.', **_option_kwds)
 @click.option('--to-ras/--no-to-ras', default=True, help='Reorient volumes to RAS before saving to TFRecords.', **_option_kwds)
 @click.option('--gzip/--no-gzip', default=True, help='Compress TFRecords with gzip (highly recommended).', **_option_kwds)
 @click.option('--verify-volumes/--no-verify-volumes', default=True, help='Verify volume pairs before converting. This option is highly recommended, as it checks that shapes of features and labels are equal to "volume-shape", that labels are (or can safely be coerced to) an integer type, and that labels are all >= 0.', **_option_kwds)
 @click.option('-j', '--num-parallel-calls', default=-1, type=int, help='Number of processes to use. If -1, uses all available processes.', **_option_kwds)
 @click.option('-v', '--verbose', is_flag=True, help='Print progress bar.', **_option_kwds)
-def convert(*, csv, tfrecords_template, volume_shape, volumes_per_shard, to_ras, gzip, verify_volumes, num_parallel_calls, verbose):
+def convert(*, csv, tfrecords_template, volume_shape, examples_per_shard, to_ras, gzip, verify_volumes, num_parallel_calls, verbose):
     """Convert medical imaging volumes to TFRecords.
 
-    Volumes must all be the same shape. This will overwrite existing TFRecords files.
+    Volumes must all be the same shape. This will overwrite existing TFRecord files.
+
+    Labels can be volumetric or scalar.
     """
     # TODO: improve docs.
     volume_filepaths = _read_csv(csv)
     num_parallel_calls = None if num_parallel_calls == -1 else num_parallel_calls
     if num_parallel_calls is None:
-        num_parallel_calls = _get_all_cpus()
+        # Get number of processes allocated to the current process.
+        # Note the difference from `os.cpu_count()`.
+        num_parallel_calls = len(os.sched_getaffinity(0))
 
     _dirname = os.path.dirname(tfrecords_template)
     if not os.path.exists(_dirname):
@@ -97,13 +99,13 @@ def convert(*, csv, tfrecords_template, volume_shape, volumes_per_shard, to_ras,
                 click.echo(pair[1])
             sys.exit(-1)
 
-    _convert(
-        volume_filepaths=volume_filepaths,
-        tfrecords_template=tfrecords_template,
-        volumes_per_shard=volumes_per_shard,
+    _write_tfrecord(
+        features_labels=volume_filepaths,
+        filename_template=tfrecords_template,
+        examples_per_shard=examples_per_shard,
         to_ras=to_ras,
-        gzip_compressed=gzip,
-        num_parallel_calls=num_parallel_calls,
+        compressed=gzip,
+        processes=num_parallel_calls,
         verbose=verbose)
 
     click.echo(click.style('Finished conversion to TFRecords.', fg='green'))
@@ -250,123 +252,6 @@ def save():
     sys.exit(-2)
 
 
-@cli.command()
-@click.option('-m', '--model', required=True, help='Neural network model architecture.', **_option_kwds)
-@click.option('-t', '--tfrecords-pattern', required=True, **_option_kwds)
-@click.option('-c', '--n-classes', type=int, required=True, **_option_kwds)
-@click.option('-b', '--batch-size', type=int, required=True, **_option_kwds)
-@click.option('-s', '--volume-shape', nargs=3, type=int, required=True, **_option_kwds)
-@click.option('--block-shape', nargs=3, type=int, help='Shape of non-overlapping blocks to take from full volumes. Default behavior is to use full volumes.', **_option_kwds)
-@click.option('-e', '--n-epochs', type=int, **_option_kwds)
-@click.option('-i', '--initial-epoch', type=int, default=0, **_option_kwds)
-@click.option('-n', '--n-volumes', type=int, required=True, help='Number of full volumes contained in the TFRecords files.', **_option_kwds)
-@click.option('-l', '--loss', required=True, help='Loss function to use.', **_option_kwds)
-@click.option('-lr', '--learning-rate', type=float, required=True, help='Learning rate for Adam optimizer.', **_option_kwds)
-@click.option('--shuffle-buffer-size', default=20, type=int, help='Size of shuffle buffer (i.e., number of volumes to shuffle at a time).', **_option_kwds)
-@click.option('--mapping', type=click.Path(exists=True), help='Path to CSV that contains replacement mapping for values in labels.', **_option_kwds)
-@click.option('--model-kwds', type=JSONParamType(), help='Keywords to pass to model, in JSON format.', **_option_kwds)
-@click.option('--augment/--no-augment', default=False, help='Apply random rigid transformations to training data. This will slow down training.', **_option_kwds)
-@click.option('--multi-gpu/--no-multi-gpu', default=False, help='Train on multiple GPUs on the same machine.', **_option_kwds)
-@click.option('-j', '--num-parallel-calls', default=-1, type=int, help='Number of processes to use for data processing. If -1, uses all available processes.', **_option_kwds)
-def train(*,
-          model,
-          tfrecords_pattern,
-          n_classes,
-          batch_size,
-          volume_shape,
-          block_shape,
-          n_epochs,
-          initial_epoch,
-          n_volumes,
-          loss,
-          learning_rate,
-          shuffle_buffer_size,
-          mapping,
-          augment,
-          model_kwds,
-          multi_gpu,
-          num_parallel_calls):
-    """Train a model."""
-
-    if num_parallel_calls == -1:
-        num_parallel_calls = _get_all_cpus()
-
-    dataset = _get_dataset(
-        file_pattern=tfrecords_pattern,
-        n_classes=n_classes,
-        batch_size=batch_size,
-        volume_shape=volume_shape,
-        block_shape=block_shape,
-        n_epochs=n_epochs,
-        mapping=mapping,
-        augment=augment,
-        shuffle_buffer_size=shuffle_buffer_size,
-        num_parallel_calls=num_parallel_calls)
-
-    steps_per_epoch = _get_steps_per_epoch(
-        n_volumes=n_volumes,
-        volume_shape=volume_shape,
-        block_shape=block_shape,
-        batch_size=batch_size)
-
-    # TODO: validation dataset
-
-    if multi_gpu:
-        # As of March 05, 2019, keras optimizers only have experimental support
-        # with MirroredStrategy training.
-        optimizer = tf.train.AdamOptimizer(learning_rate)
-    else:
-        optimizer = tf.keras.optimizers.Adam(learning_rate)
-
-    # Searches custom nobrainer losses as well as standard tf.keras losses.
-    loss = _get_loss(loss)
-    # Instantiate loss object if it is a class.
-    if inspect.isclass(loss):
-        loss = loss()
-
-
-    if model_kwds is None:
-        model_kwds = {}
-
-    click.echo('Beginning to train model in directory {}'.format(os.getcwd()))
-    click.echo('Parameters:')
-    pprint.pprint({
-        'model': model,
-        'tfrecords_pattern': tfrecords_pattern,
-        'n_classes': n_classes,
-        'batch_size': batch_size,
-        'volume_shape': volume_shape,
-        'block_shape': block_shape,
-        'n_epochs': n_epochs,
-        'initial_epoch': initial_epoch,
-        'n_volumes': n_volumes,
-        'loss': loss,
-        'learning_rate': learning_rate,
-        'shuffle_buffer_size': shuffle_buffer_size,
-        'mapping': mapping,
-        'augment': augment,
-        'model_kwds': model_kwds,
-        'multi_gpu': multi_gpu,
-        'num_parallel_calls': num_parallel_calls})
-
-    history = _train(
-        model=model,
-        dataset=dataset,
-        optimizer=optimizer,
-        loss=loss,
-        steps_per_epoch=steps_per_epoch,
-        model_kwds=model_kwds,
-        n_epochs=n_epochs,
-        initial_epoch=initial_epoch,
-        # Use default metrics and callbacks.
-        metrics=None,
-        callbacks=None,
-        multi_gpu=multi_gpu,
-        devices=None)
-
-    click.echo(click.style('Finished training model.', fg='green'))
-
-
 @cli.command()
 def evaluate():
     """Evaluate a model's predictions against known labels."""
diff --git a/nobrainer/io.py b/nobrainer/io.py
index bd07970..9eeed60 100644
--- a/nobrainer/io.py
+++ b/nobrainer/io.py
@@ -9,8 +9,6 @@ import nibabel as nib
 import numpy as np
 import tensorflow as tf
 
-from nobrainer.utils import _get_all_cpus
-
 _TFRECORDS_FEATURES_DTYPE = 'float32'
 
 
@@ -40,198 +38,17 @@ def read_mapping(filepath, skip_header=True, delimiter=','):
             "mapping values must be integers but non-integer encountered")
 
 
-def read_volume(filepath, dtype=None, return_affine=False):
+def read_volume(filepath, dtype=None, return_affine=False, to_ras=False):
     """Return numpy array of data from a neuroimaging file."""
     img = nib.load(filepath)
+    if to_ras:
+        img = nib.as_closest_canonical(img)
     data = img.get_fdata(caching='unchanged')
     if dtype is not None:
         data = data.astype(dtype)
     return data if not return_affine else (data, img.affine)
 
 
-def convert(volume_filepaths, tfrecords_template="tfrecords/data_shard-{shard:03d}.tfrecords", volumes_per_shard=100, to_ras=True, gzip_compressed=True, num_parallel_calls=None, verbose=1):
-    """Convert a list of features and labels volumes to TFRecords. The volume
-    data and the affine for each volume are saved, so the original images can
-    be recreated. This method uses multiple cores but can still take a considerable
-    amount of time for large datasets. Gzip compression should be used, as it
-    dramatically decreases the size of the resulting tfrecords file.
-
-    This method will save multiple TFRecords files if the length of `volume_filepaths` is
-    greater than `volumes_per_shard`. Sharding the data into multiple TFRecords files
-    is beneficial for at least two reasons:
-
-        1. Allows for better shuffling of large datasets because we can shuffle
-            the files and not only the underlying data.
-        2. Enables use of parallel reading of TFRecords files with the `tf.data` API,
-            which can decrease overall data processing time.
-
-    For example, if one is converting 100 pairs of files (i.e., length of
-    `volume_filepaths` is 100) and `volumes_per_shard` is 40, then three
-    TFRecords files will be created. The first 40 pairs of files will be in the
-    first shard, the next 40 pairs will be in the second shard, and the last
-    20 pairs will be in the third shard.
-
-    Parameters
-    ----------
-    volume_filepaths: nested list. Every sublist in the list should contain two
-        items: path to features volume and path to labels volume, in that order.
-    tfrecords_template: string template, path to save new tfrecords file with the string formatting
-        key 'shard'. The shard index is entered
-        Extension should be '.tfrecords'.
-    volumes_per_shard: int, number of pairs of volumes per tfrecords file.
-    to_ras: boolean, if true, reorient volumes to RAS with `nibabel.as_closest_canonical`.
-    gzip_compressed: boolean, if true, compress data with Gzip. This is highly
-        recommended, as it dramatically reduces file size.
-    num_parallel_calls: int, number of processes to use for multiprocessing. If
-        None, will use all available processes.
-    verbose: {0, 1, 2}, verbosity of the progress bar. 0 is silent, 1 is verbose,
-        and 2 is semi-verbose.
-
-    Returns
-    -------
-    None
-    """
-    try:
-        tfrecords_template.format(shard=3)
-    except Exception:
-        raise ValueError("invalid 'tfrecords_template'. This template must contain the key 'shard'.")
-
-    tfrecords_template = os.path.abspath(tfrecords_template)
-    _dirname = os.path.dirname(tfrecords_template)
-    if not os.path.exists(_dirname):
-        raise ValueError("directory does not exist: {}".format(_dirname))
-
-    n_shards = math.ceil(len(volume_filepaths) / volumes_per_shard)
-    # Include the unique tfrecords filepath for each file, because that's what
-    # the map function expects.
-    volume_filepaths_shards = [
-        [tfrecords_template.format(shard=idx), shard.tolist()]
-        for idx, shard in enumerate(np.array_split(volume_filepaths, n_shards))]
-    map_fn = functools.partial(
-        _convert, to_ras=to_ras, gzip_compressed=gzip_compressed)
-
-    print("Converting {} pairs of files to {} TFRecords.".format(len(volume_filepaths), len(volume_filepaths_shards)))
-    progbar = tf.keras.utils.Progbar(len(volume_filepaths_shards), verbose=verbose)
-    progbar.update(0)
-    if num_parallel_calls is None:
-        num_parallel_calls = _get_all_cpus()
-
-    if num_parallel_calls == 1:
-        for vf in volume_filepaths_shards:
-            map_fn(vf)
-            progbar.add(1)
-    else:
-        with multiprocessing.Pool(num_parallel_calls) as p:
-            for _ in p.imap(map_fn, volume_filepaths_shards, chunksize=2):
-                progbar.add(1)
-
-
-def _convert(tfrecords_path_volume_filepaths, to_ras=True, gzip_compressed=True):
-    """Convert a nested list of files to one TFRecords file. This function is
-    not intended for users. It is used as part of the multiprocessing function `convert`.
-
-    Parameters
-    ----------
-    tfrecords_path_volume_filepaths: list of length 2. The first item should be the path to the TFRecords file to create,
-        and the second item should be a nested list of medical imaging filepaths to convert. Each item in this nested
-        list should have two items, where the first is the path to features and the second is the path to labels.
-        This argument is structured this way to be compatible with multiprocessing map and imap methods.
-    to_ras: boolean, if true, reorient the volumes to RAS before saving to TFRecords.
-    gzip_compressed: boolean, if true, save to a Gzip-compressed file. It is highly recommended to use compression,
-        as it dramatically reduces the size of the resulting TFRecords file.
-
-    Returns
-    -------
-    None
-    """
-    tfrecords_path, volume_filepaths = tfrecords_path_volume_filepaths
-
-    def _make_one_example(features_filepath, labels_filepath):
-        """Create a `tf.train.Example` instance of the given arrays of volumetric features and labels."""
-        dtype = _TFRECORDS_FEATURES_DTYPE
-        x = nib.load(features_filepath)
-        if to_ras:
-            x = nib.as_closest_canonical(x)
-        xdata = x.get_fdata(caching='unchanged', dtype=dtype)
-        y = nib.load(labels_filepath)
-        if to_ras:
-            y = nib.as_closest_canonical(y)
-        ydata = y.get_fdata(caching='unchanged', dtype=dtype)
-
-        def bytes_feature(value):
-            return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
-
-        feature = {
-            'volume': bytes_feature(xdata.ravel().tostring()),
-            'volume_affine': bytes_feature(x.affine.astype(dtype).ravel().tostring()),
-            'label': bytes_feature(ydata.ravel().tostring()),
-            'label_affine': bytes_feature(y.affine.astype(dtype).ravel().tostring()),
-        }
-        return tf.train.Example(features=tf.train.Features(feature=feature))
-
-    if gzip_compressed:
-        options = tf.io.TFRecordOptions(compression_type="GZIP")
-    else:
-        options = tf.io.TFRecordOptions(compression_type="")
-
-    with tf.io.TFRecordWriter(tfrecords_path, options=options) as writer:
-        for fpath, lpath in volume_filepaths:
-            this_example = _make_one_example(features_filepath=fpath, labels_filepath=lpath)
-            writer.write(this_example.SerializeToString())
-
-
-def get_parse_fn(volume_shape=(256, 256, 256), include_affines=False):
-    """Return a callable that can parse examples from a TFRecords file. It is
-    assumed that each eaxmple in the TFRecords file has the keys 'volume' and
-    'label'. If 'include_affines' is true, each example must also have the keys
-    'volume_affine' and 'label_affine'.
-
-    Parameters
-    ----------
-    volume_shape: tuple of 3 ints, the shape of each volume in the TFRecords file.
-    include_affines: boolean, if true, return affine along with volumetric data.
-
-    Returns
-    -------
-    Callable used to parse a TFRecords file. This callable takes one serialized
-    example as inputs. If `include_affines` is false, this callable returns two
-    tensors: the tensor of the features volume and the tensor of the labels
-    volume. If `include_affines` if true, this callable returns four tensors:
-    the features tensor, the labels tensor, the features affine tensor, and the
-    labels affine tensor.
-    """
-    affine_shape = (4, 4)
-
-    def parse_record(serialized):
-        """Return tensors from serialized TFRecords example."""
-        features = {
-            # Would using a pre-defined shape (dependent on volume shape)
-            # improve read performance?
-            'volume': tf.io.FixedLenFeature(shape=[], dtype=tf.string),
-            'label': tf.io.FixedLenFeature(shape=[], dtype=tf.string),
-        }
-
-        if include_affines:
-            features['volume_affine'] = tf.io.FixedLenFeature(shape=[], dtype=tf.string)
-            features['label_affine'] = tf.io.FixedLenFeature(shape=[], dtype=tf.string)
-
-        def decode(s):
-            return tf.io.decode_raw(s, _TFRECORDS_FEATURES_DTYPE)
-
-        example = tf.io.parse_single_example(serialized=serialized, features=features)
-        volume = tf.reshape(decode(example['volume']), shape=volume_shape)
-        label = tf.reshape(decode(example['label']), shape=volume_shape)
-
-        if include_affines:
-            volume_affine = tf.reshape(decode(example['volume_affine']), shape=affine_shape)
-            label_affine = tf.reshape(decode(example['label_affine']), shape=affine_shape)
-            return volume, label, volume_affine, label_affine,
-        else:
-            return volume, label
-
-    return parse_record
-
-
 def verify_features_labels(volume_filepaths, volume_shape=(256, 256, 256), check_shape=True, check_labels_int=True, check_labels_gte_zero=True, num_parallel_calls=None, verbose=1):
     """Verify a list of files. This function is meant to be run before
     converting volumes to TFRecords.
@@ -239,7 +56,7 @@ def verify_features_labels(volume_filepaths, volume_shape=(256, 256, 256), check
     Parameters
     ----------
     volume_filepaths: nested list. Every sublist in the list should contain two
-        items: path to features volume and path to labels volume, in that order.
+        items: (1) path to feature volume and (2) path to label volume or a scalar.
     volume_shape: tuple of three ints. Shape that both volumes should be.
     check_shape: boolean, if true, validate that the shape of both volumes is
         equal to 'volume_shape'.
@@ -257,23 +74,36 @@ def verify_features_labels(volume_filepaths, volume_shape=(256, 256, 256), check
     List of invalid pairs of filepaths. If the list is empty, all filepaths are
     valid.
     """
+    from nobrainer.tfrecord import _labels_all_scalar
 
     for pair in volume_filepaths:
         if len(pair) != 2:
             raise ValueError(
                 "all items in 'volume_filepaths' must have length of 2, but"
                 " found at least one item with lenght != 2.")
+
+    labels = (y for _, y in volume_filepaths)
+    scalar_labels = _labels_all_scalar(labels)
+
+    for pair in volume_filepaths:
         if not os.path.exists(pair[0]):
             raise ValueError("file does not exist: {}".format(pair[0]))
-        if not os.path.exists(pair[1]):
-            raise ValueError("file does not exist: {}".format(pair[1]))
+        if not scalar_labels:
+            if not os.path.exists(pair[1]):
+                raise ValueError("file does not exist: {}".format(pair[1]))
 
-    print("Verifying {} pairs of volumes".format(len(volume_filepaths)))
+    if scalar_labels:
+        map_fn = functools.partial(_verify_features_scalar_labels, volume_shape=volume_shape)
+    else:
+        map_fn = functools.partial(_verify_features_nonscalar_labels, volume_shape=volume_shape, check_shape=check_shape, check_labels_int=check_labels_int, check_labels_gte_zero=check_labels_gte_zero)
+    if num_parallel_calls is None:
+        # Get number of processes allocated to the current process.
+        # Note the difference from `os.cpu_count()`.
+        num_parallel_calls = len(os.sched_getaffinity(0))
+
+    print("Verifying {} examples".format(len(volume_filepaths)))
     progbar = tf.keras.utils.Progbar(len(volume_filepaths), verbose=verbose)
     progbar.update(0)
-    map_fn = functools.partial(_verify_features_labels_pair, volume_shape=volume_shape, check_shape=check_shape, check_labels_int=check_labels_int, check_labels_gte_zero=check_labels_gte_zero)
-    if num_parallel_calls is None:
-        num_parallel_calls = _get_all_cpus()
 
     outputs = []
     if num_parallel_calls == 1:
@@ -290,7 +120,7 @@ def verify_features_labels(volume_filepaths, volume_shape=(256, 256, 256), check
     return invalid_files
 
 
-def _verify_features_labels_pair(pair_of_paths, *, volume_shape, check_shape, check_labels_int, check_labels_gte_zero):
+def _verify_features_nonscalar_labels(pair_of_paths, *, volume_shape, check_shape, check_labels_int, check_labels_gte_zero):
     """Verify a pair of features and labels volumes."""
     x = nib.load(pair_of_paths[0])
     y = nib.load(pair_of_paths[1])
@@ -315,6 +145,19 @@ def _verify_features_labels_pair(pair_of_paths, *, volume_shape, check_shape, ch
     return True
 
 
+def _verify_features_scalar_labels(path_scalar, *, volume_shape):
+    """Check that feature has the desired shape and that label is scalar."""
+    from nobrainer.tfrecord import _is_int_or_float
+
+    feature, label = path_scalar
+    x = nib.load(feature)
+    if x.shape != volume_shape:
+        return False
+    if not _is_int_or_float(label):
+        return False
+    return True
+
+
 def _is_gzipped(filepath):
     """Return True if the file is gzip-compressed, False otherwise."""
     with open(filepath, 'rb') as f:
diff --git a/nobrainer/tfrecord.py b/nobrainer/tfrecord.py
new file mode 100644
index 0000000..7a07ed7
--- /dev/null
+++ b/nobrainer/tfrecord.py
@@ -0,0 +1,274 @@
+import functools
+import math
+import multiprocessing as mp
+import os
+from pathlib import Path
+import warnings
+
+import numpy as np
+import tensorflow as tf
+
+from nobrainer.io import read_volume
+
+_TFRECORDS_DTYPE = "float32"
+
+
+def write(features_labels, filename_template, examples_per_shard,
+        to_ras=True,
+        compressed=True, processes=None, chunksize=1, verbose=1):
+    """Write to TFRecords files."""
+    n_examples = len(features_labels)
+    n_shards = math.ceil(n_examples / examples_per_shard)
+    shards = np.array_split(features_labels, n_shards)
+
+    # Test that the `filename_template` has a `shard` formatting key.
+    try:
+        filename_template.format(shard=0)
+    except Exception:
+        raise ValueError("`filename_template` must include a string formatting key 'shard' that accepts an integer.")
+
+    # This is the object that returns a protocol buffer string of the feature and label on each iteration.
+    # It is pickle-able, unlike a generator.
+    proto_iterators = [_ProtoIterator(s) for s in shards]
+    # Set up positional arguments for the core writer function.
+    iterable = [(p, filename_template.format(shard=j)) for j, p in enumerate(proto_iterators)]
+    # Set keyword arguments so the resulting function accepts one positional argument.
+    map_fn = functools.partial(_write_tfrecords, compressed=True)
+
+    # This is a hack to allow multiprocessing to pickle
+    # the __writer_func object. Pickles don't like local functions.
+    global __writer_func
+
+    def __writer_func(iterator_filename):
+        iterator, filename = iterator_filename
+        map_fn(protobuf_iterator=iterator, filename=filename)
+
+    progbar = tf.keras.utils.Progbar(target=len(iterable), verbose=verbose)
+    progbar.update(0)
+    if processes is None:
+        # Get number of processes allocated to the current process.
+        # Note the difference from `os.cpu_count()`.
+        processes = len(os.sched_getaffinity(0))
+    with mp.Pool(processes) as p:
+        for _ in p.imap_unordered(__writer_func, iterable=iterable, chunksize=chunksize):
+            progbar.add(1)
+
+
+def parse_example_fn(volume_shape, scalar_label=False):
+    """"""
+    @tf.function
+    def parse_example(serialized):
+        """"""
+        features = {
+            "feature/shape": tf.io.FixedLenFeature(shape=[], dtype=tf.string),
+            "feature/value": tf.io.FixedLenFeature(shape=[], dtype=tf.string),
+            "label/value": tf.io.FixedLenFeature(shape=[], dtype=tf.string),
+            "label/rank": tf.io.FixedLenFeature(shape=[], dtype=tf.int64)}
+        e = tf.io.parse_single_example(serialized=serialized, features=features)
+        x = tf.io.decode_raw(e["feature/value"], _TFRECORDS_DTYPE)
+        y = tf.io.decode_raw(e["label/value"], _TFRECORDS_DTYPE)
+        # TODO: this line does not work. The shape cannot be determined
+        # dynamically... for now.
+        # xshape = tf.cast(tf.io.decode_raw(e["feature/shape"], _TFRECORDS_DTYPE), tf.int32)
+        x = tf.reshape(x, shape=volume_shape)
+        if not scalar_label:
+            y = tf.reshape(y, shape=volume_shape)
+        else:
+            y = tf.reshape(y, shape=[1])
+        return x, y
+    return parse_example
+
+
+def _bytes_feature(value):
+    """Returns a bytes_list from a string / byte."""
+    if isinstance(value, type(tf.constant(0))):
+        value = value.numpy() # BytesList won't unpack a string from an EagerTensor.
+    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
+
+
+def _float_feature(value):
+    """Returns a float_list from a float / double."""
+    return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))
+
+
+def _int64_feature(value):
+    """Returns an int64_list from a bool / enum / int / uint."""
+    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
+
+
+def _dtype_to_bytes(dtype):
+    if isinstance(dtype, str):
+        return dtype.encode('utf-8')
+    try:
+        s = str(dtype.name)
+    except AttributeError:
+        s = str(dtype)
+    return s.encode('utf-8')
+
+
+def _get_feature_dict(x, affine):
+    x = np.asarray(x)
+    x = x.astype(_TFRECORDS_DTYPE)
+
+    feature = {
+        "feature/value": _bytes_feature(x.tobytes()),
+        "feature/dtype": _bytes_feature(_dtype_to_bytes(x.dtype)),
+        "feature/rank": _int64_feature(x.ndim)}
+
+    for j, s in enumerate(x.shape):
+        feature["feature/shape/dim{}".format(j)] = _int64_feature(s)
+
+    if x.ndim:
+        feature["feature/shape"] = _bytes_feature(np.array(x.shape).astype(_TFRECORDS_DTYPE).tobytes())
+
+    affine = np.asarray(affine)
+    affine = affine.astype(_TFRECORDS_DTYPE)
+    feature.update({
+        "feature/affine/value": _bytes_feature(affine.tobytes()),
+        "feature/affine/dtype": _bytes_feature(_dtype_to_bytes(affine.dtype)),
+        "feature/affine/rank": _int64_feature(affine.ndim),
+    })
+    for j, s in enumerate(affine.shape):
+        feature["feature/affine/shape/dim{}".format(j)] = _int64_feature(s)
+
+    return feature
+
+
+def _get_label_dict(y, affine=None):
+    y = np.asarray(y)
+    y = y.astype(_TFRECORDS_DTYPE)
+
+    label = {
+        "label/value": _bytes_feature(y.tobytes()),
+        "label/dtype": _bytes_feature(_dtype_to_bytes(y.dtype)),
+        "label/rank": _int64_feature(y.ndim)}
+
+    for j, s in enumerate(y.shape):
+        label["label/shape/dim{}".format(j)] = _int64_feature(s)
+    if y.ndim:
+        label["feature/shape"] = _bytes_feature(np.array(y.shape).astype(_TFRECORDS_DTYPE).tobytes())
+
+    if affine is None and y.ndim != 0:
+        raise ValueError("Affine is required when label is not scalar.")
+
+    if affine is not None:
+        affine = np.asarray(affine)
+        affine = affine.astype(_TFRECORDS_DTYPE)
+        label.update({
+            "label/affine/value": _bytes_feature(affine.tobytes()),
+            "label/affine/dtype": _bytes_feature(_dtype_to_bytes(affine.dtype)),
+            "label/affine/rank": _int64_feature(affine.ndim),
+        })
+        for j, s in enumerate(affine.shape):
+            label["label/affine/shape/dim{}".format(j)] = _int64_feature(s)
+
+    return label
+
+
+def _to_proto(feature, label, feature_affine, label_affine=None):
+    """Return instance of `tf.train.Example` of feature and label."""
+    d = _get_feature_dict(x=feature, affine=feature_affine)
+    d.update(_get_label_dict(y=label, affine=label_affine))
+    return tf.train.Example(features=tf.train.Features(feature=d))
+
+
+class _ProtoIterator:
+    """Iterator of protobuf strings.
+
+    Parameters
+    ----------
+
+    This custom iterator is used instead of a generator for use in
+    multiprocessing. Generators cannot be pickled, and so cannot be used in
+    multiprocessing workflows. Please see
+    https://stackoverflow.com/a/7180424/5666087 for more information.
+    """
+    def __init__(self, features_labels, to_ras=True):
+        self.features_labels = features_labels
+        self.to_ras = to_ras
+
+        # Try to "intelligently" deduce if the labels are scalars or not.
+        # An alternative here would be to check if these point to existing
+        # files, though it is possible to have existing filenames that
+        # are integers or floats.
+        labels = [y for _, y in features_labels]
+        self.scalar_label = _labels_all_scalar(labels)
+        self._j = 0
+
+        no_exist = []
+        for x, _ in self.features_labels:
+            if not Path(x).exists():
+                no_exist.append(x)
+        if no_exist:
+            raise ValueError("Some files do not exist: {}".format(", ".join(no_exist)))
+
+        if not self.scalar_label:
+            no_exist = []
+            for _, y in self.features_labels:
+                if not Path(y).exists():
+                    no_exist.append(y)
+            if no_exist:
+                raise ValueError("Some files do not exist: {}".format(", ".join(no_exist)))
+
+    def __iter__(self):
+        self._j = 0
+        return self
+
+    def __next__(self):
+        index = self._j
+        serialized = self._serialize(index)
+        self._j += 1
+        return serialized
+
+    def _serialize(self, index):
+        try:
+            x, y = self.features_labels[index]
+        except IndexError:
+            raise StopIteration
+        if self.scalar_label:
+            x, affine = read_volume(x, return_affine=True, dtype=_TFRECORDS_DTYPE, to_ras=self.to_ras)
+            proto = _to_proto(feature=x, label=y, feature_affine=affine, label_affine=None)
+            return proto.SerializeToString()
+        else:
+            x, affine_x = read_volume(x, return_affine=True, dtype=_TFRECORDS_DTYPE, to_ras=self.to_ras)
+            y, affine_y = read_volume(y, return_affine=True, dtype=_TFRECORDS_DTYPE, to_ras=self.to_ras)
+            proto = _to_proto(feature=x, label=y, feature_affine=affine_x, label_affine=affine_y)
+            return proto.SerializeToString()
+
+
+def _write_tfrecords(protobuf_iterator, filename, compressed=True):
+    """
+    protobuf_iterator: iterator, iterator which yields protocol-buffer serialized strings.
+    """
+    if compressed:
+        options = tf.io.TFRecordOptions(compression_type="GZIP")
+    else:
+        options = None
+    with tf.io.TFRecordWriter(filename, options=options) as f:
+        for proto_string in protobuf_iterator:
+            f.write(proto_string)
+
+
+def _is_int_or_float(value):
+    if isinstance(value, (int, float)):
+        return True
+    try:
+        float(value)
+        return True
+    except (TypeError, ValueError):
+        pass
+    try:
+        int(value)
+        return True
+    except (TypeError, ValueError):
+        pass
+    return False
+
+
+def _labels_all_scalar(labels):
+    scalars = list(map(_is_int_or_float, labels))
+    if any(scalars) and not all(scalars):
+        raise ValueError(
+            "Some labels were detected as scalars, while others were not."
+            " Labels must be all scalars or all filenames of volumes.")
+    return all(scalars)
diff --git a/nobrainer/training.py b/nobrainer/training.py
deleted file mode 100644
index 8ee59e9..0000000
--- a/nobrainer/training.py
+++ /dev/null
@@ -1,128 +0,0 @@
-"""Methods for training models."""
-
-import os
-
-import tensorflow as tf
-
-from nobrainer import metrics
-from nobrainer import models
-
-
-def train(model, dataset, optimizer, loss, steps_per_epoch, model_kwds={}, n_epochs=1, initial_epoch=0, metrics=None, callbacks=None, validation_dataset=None, validation_steps=None, multi_gpu=False, devices=None):
-    """Train a model.
-
-    """
-
-    # TODO: add checks.
-    n_classes = dataset.output_shapes[1][-1]  # last dimension of labels
-    input_shape = dataset.output_shapes[0][1:]  # remove batch dimension
-
-    if metrics is None:
-        metrics = _get_default_metrics(n_classes)
-
-    if multi_gpu:
-        model = _compile_multi_gpu_model(
-            model=model,
-            n_classes=n_classes,
-            input_shape=input_shape,
-            optimizer=optimizer,
-            loss=loss,
-            model_kwds=model_kwds,
-            metrics=metrics,
-            devices=devices)
-    else:
-        if isinstance(model, str):
-            model_fn = models.get(model)
-            model = model_fn(n_classes=n_classes, input_shape=input_shape, **model_kwds)
-            model.compile(optimizer=optimizer, loss=loss, metrics=metrics)
-        elif isinstance(mode, tf.keras.Model):
-            # TODO: can we test if the model is compiled? We lose the optimizer
-            # state if it was already compiled.
-            model.compile(optimizer=optimizer, loss=loss, metrics=metrics)
-
-    history = _train_generic(
-        compiled_model=model,
-        dataset=dataset,
-        steps_per_epoch=steps_per_epoch,
-        n_epochs=n_epochs,
-        initial_epoch=initial_epoch,
-        callbacks=callbacks,
-        validation_dataset=validation_dataset,
-        validation_steps=validation_steps)
-
-    return history
-
-
-def _train_generic(compiled_model, dataset, steps_per_epoch, n_epochs, initial_epoch, callbacks=None, validation_dataset=None, validation_steps=None):
-    """"""
-
-    modelcheckpoint_path = './checkpoints/model-{}_epoch-{{epoch:03d}}.ckpt'.format(compiled_model.name)
-    if callbacks is None:
-        tensorboard_dir = './logs'
-        os.makedirs(tensorboard_dir, exist_ok=True)
-        os.makedirs(os.path.dirname(modelcheckpoint_path), exist_ok=True)
-        callbacks = [
-            tf.keras.callbacks.ModelCheckpoint(
-                modelcheckpoint_path, save_weights_only=True),
-            # Write metrics out every 100 samples. Writing too frequently can
-            # slow down training.
-            tf.keras.callbacks.TensorBoard(
-                tensorboard_dir, write_graph=True, update_freq=100),
-        ]
-
-    # TODO: if we can load weights after compiling model, load the most recent
-    # checkpoint before training.
-    try:
-        compiled_model.load_weights(modelcheckpoint_path)
-    except tf.errors.NotFoundError:
-        pass
-
-    history = compiled_model.fit(
-        dataset,
-        epochs=n_epochs,
-        initial_epoch=initial_epoch,
-        steps_per_epoch=steps_per_epoch,
-        validation_data=validation_dataset,
-        validation_steps=validation_steps,
-        callbacks=callbacks)
-
-    return history
-
-
-def _compile_multi_gpu_model(model, n_classes, input_shape, optimizer, loss, model_kwds={}, metrics=None, devices=None):
-    """Train a model across multiple GPUs on the same machine."""
-
-    if isinstance(model, tf.keras.Model):
-        raise ValueError(
-            "Model cannot be a `tf.keras.Model` object, because it must be"
-            " instantiated within the MirroredStrategy scope.")
-
-    model_fn = models.get(model)
-
-    strategy = tf.distribute.MirroredStrategy()
-    with strategy.scope():
-        model = model_fn(n_classes=n_classes, input_shape=input_shape, **model_kwds)
-        model.compile(optimizer=optimizer, loss=loss, metrics=metrics)
-
-    return model
-
-
-def _train_tpu():
-    """"""
-    raise NotImplementedError()
-
-
-def _get_default_metrics(n_classes):
-    if n_classes == 1:
-        m = [
-            metrics.dice,
-        ]
-    elif n_classes > 2:
-        m = [
-            metrics.generalized_dice,
-            metrics.tversky,
-        ]
-    else:
-        raise ValueError("Invalid number of classes. Got {}".format(n_classes))
-
-    return m
diff --git a/nobrainer/utils.py b/nobrainer/utils.py
index 3192174..4a9b7d8 100644
--- a/nobrainer/utils.py
+++ b/nobrainer/utils.py
@@ -95,15 +95,3 @@ def get_data(cache_dir=_cache_dir):
         writer.writerows(output)
 
     return csvpath
-
-
-def _get_all_cpus():
-    """SLURM-aware function that returns number of CPUs or None if not using
-    SLURM.
-    """
-    nproc = os.environ.get('SLURM_CPUS_ON_NODE', None)
-    if nproc is not None:
-        nproc = int(nproc)
-    else:
-        nproc = multiprocessing.cpu_count()
-    return nproc
diff --git a/nobrainer/volume.py b/nobrainer/volume.py
index 346090b..087b9d3 100644
--- a/nobrainer/volume.py
+++ b/nobrainer/volume.py
@@ -8,7 +8,7 @@ import numpy as np
 import tensorflow as tf
 
 from nobrainer.io import _is_gzipped
-from nobrainer.io import get_parse_fn
+from nobrainer.tfrecord import parse_example_fn
 from nobrainer.transform import get_affine
 from nobrainer.transform import warp_features_labels
 
@@ -78,7 +78,7 @@ def get_dataset(file_pattern, n_classes, batch_size, volume_shape,
 
     # Parse each example in each TFRecords file as a tensor of features and a
     # tensor of labels.
-    parse_fn = get_parse_fn(volume_shape=volume_shape)
+    parse_fn = parse_example_fn(volume_shape=volume_shape, scalar_label=False)
     dataset = dataset.map(map_func=parse_fn)
 
     # At this point, dataset output will be two tensors, both with shape


Test patch:
diff --git a/nobrainer/cli/tests/main_test.py b/nobrainer/cli/tests/main_test.py
index 1d94185..c90b3e9 100644
--- a/nobrainer/cli/tests/main_test.py
+++ b/nobrainer/cli/tests/main_test.py
@@ -1,5 +1,6 @@
 """Tests for `nobrainer.cli.main`."""
 
+import csv
 from pathlib import Path
 
 import click
@@ -9,14 +10,14 @@ import numpy as np
 import pytest
 
 from nobrainer.cli import main as climain
-from nobrainer.io import convert
 from nobrainer.io import read_csv
 from nobrainer.io import read_volume
 from nobrainer.models.meshnet import meshnet
+from nobrainer.tfrecord import write
 from nobrainer.utils import get_data
 
 
-def test_convert(tmp_path):
+def test_convert_nonscalar_labels(tmp_path):
     runner = CliRunner()
     with runner.isolated_filesystem():
         csvpath = get_data(str(tmp_path))
@@ -24,7 +25,59 @@ def test_convert(tmp_path):
         tfrecords_template.parent.mkdir(exist_ok=True)
         args = """\
     convert --csv={} --tfrecords-template={} --volume-shape 256 256 256
-        --volumes-per-shard=2 --to-ras --no-verify-volumes
+        --examples-per-shard=2 --to-ras --no-verify-volumes
+    """.format(csvpath, tfrecords_template)
+        result = runner.invoke(climain.cli, args.split())
+        assert result.exit_code == 0
+        assert Path('data/shard-000.tfrecords').is_file()
+        assert Path('data/shard-001.tfrecords').is_file()
+        assert Path('data/shard-002.tfrecords').is_file()
+        assert Path('data/shard-003.tfrecords').is_file()
+        assert Path('data/shard-004.tfrecords').is_file()
+        assert not Path('data/shard-005.tfrecords').is_file()
+
+
+def test_convert_scalar_int_labels(tmp_path):
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        csvpath = get_data(str(tmp_path))
+        # Make labels scalars.
+        data = [(x, 0) for (x, _) in read_csv(csvpath)]
+        csvpath = tmp_path.with_suffix(".new.csv")
+        with open(csvpath, 'w', newline='') as myfile:
+            wr = csv.writer(myfile, quoting=csv.QUOTE_ALL)
+            wr.writerows(data)
+        tfrecords_template = Path('data/shard-{shard:03d}.tfrecords')
+        tfrecords_template.parent.mkdir(exist_ok=True)
+        args = """\
+    convert --csv={} --tfrecords-template={} --volume-shape 256 256 256
+        --examples-per-shard=2 --to-ras --no-verify-volumes
+    """.format(csvpath, tfrecords_template)
+        result = runner.invoke(climain.cli, args.split())
+        assert result.exit_code == 0
+        assert Path('data/shard-000.tfrecords').is_file()
+        assert Path('data/shard-001.tfrecords').is_file()
+        assert Path('data/shard-002.tfrecords').is_file()
+        assert Path('data/shard-003.tfrecords').is_file()
+        assert Path('data/shard-004.tfrecords').is_file()
+        assert not Path('data/shard-005.tfrecords').is_file()
+
+
+def test_convert_scalar_int_labels(tmp_path):
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        csvpath = get_data(str(tmp_path))
+        # Make labels scalars.
+        data = [(x, 1.0) for (x, _) in read_csv(csvpath)]
+        csvpath = tmp_path.with_suffix(".new.csv")
+        with open(csvpath, 'w', newline='') as myfile:
+            wr = csv.writer(myfile, quoting=csv.QUOTE_ALL)
+            wr.writerows(data)
+        tfrecords_template = Path('data/shard-{shard:03d}.tfrecords')
+        tfrecords_template.parent.mkdir(exist_ok=True)
+        args = """\
+    convert --csv={} --tfrecords-template={} --volume-shape 256 256 256
+        --examples-per-shard=2 --to-ras --no-verify-volumes
     """.format(csvpath, tfrecords_template)
         result = runner.invoke(climain.cli, args.split())
         assert result.exit_code == 0
@@ -68,29 +121,6 @@ def test_save():
     assert False
 
 
-def test_train():
-    runner = CliRunner()
-    with runner.isolated_filesystem():
-        xpath = 'features.nii.gz'
-        ypath = 'labels.nii.gz'
-        nib.Nifti1Image(np.ones((20, 20, 20)), np.eye(4)).to_filename(str(xpath))
-        nib.Nifti1Image(np.ones((20, 20, 20)), np.eye(4)).to_filename(str(ypath))
-        files = [(str(xpath), str(ypath))]
-        convert(files, tfrecords_template='data-{shard:03d}.tf')
-
-        args = """\
-    train --model=meshnet --tfrecords-pattern={} --n-classes=1 --batch-size=1
-        --volume-shape 20 20 20 --block-shape 10 10 10 --n-epochs=1 --n-volumes=1
-        --loss=dice --learning-rate=1e-05
-    """.format('data-*.tf')
-
-        result = runner.invoke(climain.cli, args.split())
-        print(result.output)
-        assert result.exit_code == 0
-        assert Path('logs').is_dir()
-        assert Path('checkpoints').is_dir()
-
-
 @pytest.mark.xfail
 def test_evaluate():
     assert False
diff --git a/nobrainer/models/tests/models_test.py b/nobrainer/models/tests/models_test.py
index 062133a..7c70854 100644
--- a/nobrainer/models/tests/models_test.py
+++ b/nobrainer/models/tests/models_test.py
@@ -8,6 +8,7 @@ from nobrainer.models.meshnet import meshnet_vwn
 from nobrainer.models.unet import unet
 from nobrainer.models.autoencoder import autoencoder
 
+
 def model_test(model_cls, n_classes, input_shape, kwds={}):
     """Tests for models."""
     x = 10 * np.random.random(input_shape)
diff --git a/nobrainer/tests/io_test.py b/nobrainer/tests/io_test.py
index 9ffd66b..d742ca3 100644
--- a/nobrainer/tests/io_test.py
+++ b/nobrainer/tests/io_test.py
@@ -71,42 +71,47 @@ def test_read_volume(tmp_path):
     assert np.array_equal(data, data_loaded)
     assert np.array_equal(affine, affine_loaded)
 
+    data = np.random.rand(8, 8, 8).astype(np.float32)
+    affine = np.array([
+        [1.5, 0, 1.2, 0],
+        [0.8, 0.8, 0, 0],
+        [0, 1, 0, 0],
+        [0, 0, 0, 1]
+    ])
+    filename = str(tmp_path / 'foo_asr.nii.gz')
+    nib.save(nib.Nifti1Image(data, affine), filename)
+    data_loaded = io.read_volume(filename, to_ras=True)
+    assert not np.array_equal(data, data_loaded)
+    data_loaded = io.read_volume(filename, to_ras=False)
+    assert np.array_equal(data, data_loaded)
+
 
-def test_convert(csv_of_volumes, tmp_path):
+def test_verify_features_nonscalar_labels(csv_of_volumes):
     files = io.read_csv(csv_of_volumes, skip_header=False)
-    tfrecords_template = str(tmp_path / 'data-{shard:03d}.tfrecords')
-    volumes_per_shard = 12
-    io.convert(
-        files,
-        tfrecords_template=tfrecords_template,
-        volumes_per_shard=volumes_per_shard,
-        num_parallel_calls=1)
-
-    paths = list(tmp_path.glob('data-*.tfrecords'))
-    paths = sorted(paths)
-    assert len(paths) == 9
-    assert (tmp_path / 'data-008.tfrecords').is_file()
-
-    dset = tf.data.TFRecordDataset(list(map(str, paths)), compression_type='GZIP')
-    dset = dset.map(io.get_parse_fn(volume_shape=(8, 8, 8), include_affines=True))
-
-    for ref, test in zip(files, dset):
-        x, y = ref
-        x, x_aff = io.read_volume(x, return_affine=True)
-        y, y_aff = io.read_volume(y, return_affine=True)
-        assert np.array_equal(x, test[0])
-        assert np.array_equal(y, test[1])
-        assert np.array_equal(x_aff, test[2])
-        assert np.array_equal(y_aff, test[3])
-
-    with pytest.raises(ValueError):
-        io.convert(files, tfrecords_template="data/foobar-{}.tfrecords")
-
-
-def test_verify_features_labels(csv_of_volumes):
+    invalid = io.verify_features_labels(
+        files, volume_shape=(8, 8, 8), num_parallel_calls=1)
+    assert not invalid
+    # TODO: add more cases.
+
+
+def test_verify_features_scalar_labels(csv_of_volumes):
     files = io.read_csv(csv_of_volumes, skip_header=False)
-    io.verify_features_labels(
+    # Int labels.
+    files = [(x, 0) for (x, _) in files]
+    invalid = io.verify_features_labels(
+        files, volume_shape=(8, 8, 8), num_parallel_calls=1)
+    assert not invalid
+    invalid = io.verify_features_labels(
+        files, volume_shape=(12, 12, 8), num_parallel_calls=1)
+    assert all(invalid)
+    # Float labels.
+    files = [(x, 1.0) for (x, _) in files]
+    invalid = io.verify_features_labels(
         files, volume_shape=(8, 8, 8), num_parallel_calls=1)
+    assert not invalid
+    invalid = io.verify_features_labels(
+        files, volume_shape=(12, 12, 8), num_parallel_calls=1)
+    assert all(invalid)
 
 
 def test_is_gzipped(tmp_path):
diff --git a/nobrainer/tests/test_utils.py b/nobrainer/tests/test_utils.py
index c19061e..1125448 100644
--- a/nobrainer/tests/test_utils.py
+++ b/nobrainer/tests/test_utils.py
@@ -16,9 +16,3 @@ def test_get_data():
     for x, y in files:
         assert Path(x).is_file()
         assert Path(y).is_file()
-
-
-def test_get_all_cpus():
-    assert nbutils._get_all_cpus() == multiprocessing.cpu_count()
-    os.environ['SLURM_CPUS_ON_NODE'] = "128"
-    assert nbutils._get_all_cpus() == 128
diff --git a/nobrainer/tests/tfrecord_test.py b/nobrainer/tests/tfrecord_test.py
new file mode 100644
index 0000000..445d897
--- /dev/null
+++ b/nobrainer/tests/tfrecord_test.py
@@ -0,0 +1,163 @@
+import random
+
+import numpy as np
+from numpy.testing import assert_array_equal
+import pytest
+import tensorflow as tf
+
+from nobrainer import io
+from nobrainer import tfrecord
+from nobrainer.tests.utils import csv_of_volumes
+
+
+def test_write_read_volume_labels(csv_of_volumes, tmp_path):
+    files = io.read_csv(csv_of_volumes, skip_header=False)
+    filename_template = str(tmp_path / 'data-{shard:03d}.tfrecords')
+    examples_per_shard = 12
+    tfrecord.write(
+        files,
+        filename_template=filename_template,
+        examples_per_shard=examples_per_shard,
+        processes=1,
+    )
+
+    paths = list(tmp_path.glob('data-*.tfrecords'))
+    paths = sorted(paths)
+    assert len(paths) == 9
+    assert (tmp_path / 'data-008.tfrecords').is_file()
+
+    dset = tf.data.TFRecordDataset(list(map(str, paths)), compression_type='GZIP')
+    dset = dset.map(tfrecord.parse_example_fn(volume_shape=(8, 8, 8), scalar_label=False))
+
+    for ref, test in zip(files, dset):
+        x, y = ref
+        x, y = io.read_volume(x), io.read_volume(y)
+        assert_array_equal(x, test[0])
+        assert_array_equal(y, test[1])
+
+    with pytest.raises(ValueError):
+        tfrecord.write(
+            files, filename_template="data/foobar-{}.tfrecords",
+            examples_per_shard=4
+        )
+
+
+def test_write_read_volume_labels_all_processes(csv_of_volumes, tmp_path):
+    files = io.read_csv(csv_of_volumes, skip_header=False)
+    filename_template = str(tmp_path / 'data-{shard:03d}.tfrecords')
+    examples_per_shard = 12
+    tfrecord.write(
+        files,
+        filename_template=filename_template,
+        examples_per_shard=examples_per_shard,
+        processes=None,
+    )
+
+    paths = list(tmp_path.glob('data-*.tfrecords'))
+    paths = sorted(paths)
+    assert len(paths) == 9
+    assert (tmp_path / 'data-008.tfrecords').is_file()
+
+    dset = tf.data.TFRecordDataset(list(map(str, paths)), compression_type='GZIP')
+    dset = dset.map(tfrecord.parse_example_fn(volume_shape=(8, 8, 8), scalar_label=False))
+
+    for ref, test in zip(files, dset):
+        x, y = ref
+        x, y = io.read_volume(x), io.read_volume(y)
+        assert_array_equal(x, test[0])
+        assert_array_equal(y, test[1])
+
+    with pytest.raises(ValueError):
+        tfrecord.write(
+            files, filename_template="data/foobar-{}.tfrecords",
+            examples_per_shard=4
+        )
+
+
+def test_write_read_float_labels(csv_of_volumes, tmp_path):
+    files = io.read_csv(csv_of_volumes, skip_header=False)
+    files = [(x, random.random()) for x, _ in files]
+    filename_template = str(tmp_path / 'data-{shard:03d}.tfrecords')
+    examples_per_shard = 12
+    tfrecord.write(
+        files,
+        filename_template=filename_template,
+        examples_per_shard=examples_per_shard,
+        processes=1,
+    )
+
+    paths = list(tmp_path.glob('data-*.tfrecords'))
+    paths = sorted(paths)
+    assert len(paths) == 9
+    assert (tmp_path / 'data-008.tfrecords').is_file()
+
+    dset = tf.data.TFRecordDataset(list(map(str, paths)), compression_type='GZIP')
+    dset = dset.map(tfrecord.parse_example_fn(volume_shape=(8, 8, 8), scalar_label=True))
+
+    for ref, test in zip(files, dset):
+        x, y = ref
+        x = io.read_volume(x)
+        assert_array_equal(x, test[0])
+        assert_array_equal(y, test[1])
+
+
+def test_write_read_int_labels(csv_of_volumes, tmp_path):
+    files = io.read_csv(csv_of_volumes, skip_header=False)
+    files = [(x, random.randint(0, 9)) for x, _ in files]
+    filename_template = str(tmp_path / 'data-{shard:03d}.tfrecords')
+    examples_per_shard = 12
+    tfrecord.write(
+        files,
+        filename_template=filename_template,
+        examples_per_shard=examples_per_shard,
+        processes=1)
+
+    paths = list(tmp_path.glob('data-*.tfrecords'))
+    paths = sorted(paths)
+    assert len(paths) == 9
+    assert (tmp_path / 'data-008.tfrecords').is_file()
+
+    dset = tf.data.TFRecordDataset(list(map(str, paths)), compression_type='GZIP')
+    dset = dset.map(tfrecord.parse_example_fn(volume_shape=(8, 8, 8), scalar_label=True))
+
+    for ref, test in zip(files, dset):
+        x, y = ref
+        x = io.read_volume(x)
+        assert_array_equal(x, test[0])
+        assert_array_equal(y, test[1])
+
+
+def test__is_int_or_float():
+    assert tfrecord._is_int_or_float(10)
+    assert tfrecord._is_int_or_float(10.0)
+    assert tfrecord._is_int_or_float("10")
+    assert tfrecord._is_int_or_float("10.00")
+    assert not tfrecord._is_int_or_float("foobar")
+    assert tfrecord._is_int_or_float(np.ones(1))
+    assert not tfrecord._is_int_or_float(np.ones(10))
+
+
+def test__dtype_to_bytes():
+    np_tf_dt = [
+        (np.uint8, tf.uint8, b'uint8'),
+        (np.uint16, tf.uint16, b'uint16'),
+        (np.uint32, tf.uint32, b'uint32'),
+        (np.uint64, tf.uint64, b'uint64'),
+
+        (np.int8, tf.int8, b'int8'),
+        (np.int16, tf.int16, b'int16'),
+        (np.int32, tf.int32, b'int32'),
+        (np.int64, tf.int64, b'int64'),
+
+        (np.float16, tf.float16, b'float16'),
+        (np.float32, tf.float32, b'float32'),
+        (np.float64, tf.float64, b'float64'),
+    ]
+
+    for npd, tfd, dt in np_tf_dt:
+        npd = np.dtype(npd)
+        assert tfrecord._dtype_to_bytes(npd) == dt
+        assert tfrecord._dtype_to_bytes(tfd) == dt
+
+    assert tfrecord._dtype_to_bytes('float32') == b'float32'
+    assert tfrecord._dtype_to_bytes('foobar') == b'foobar'