notebooks/ML4H_Model_Factory_Intro.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to ML using ML4H"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "- Basic comfort with python, some linear algebra, some data science\n",
    "- Follow the instructions in the main [README](https://github.com/broadinstitute/ml4h) for installing ML4H\n",
    "- Data used in this notebook is available here ([gs://fc-500bd872-4a53-45c9-87d3-39656bd83f85/data/hg002_na24385_ml4h_tensors_v2021_10_14.tar.gz](gs://fc-500bd872-4a53-45c9-87d3-39656bd83f85/data/hg002_na24385_ml4h_tensors_v2021_10_14.tar.gz))\n",
    "- Now we are ready to teach the machines!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports\n",
    "import os\n",
    "import sys\n",
    "import pickle\n",
    "import gzip\n",
    "from typing import Dict\n",
    "\n",
    "import h5py\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "from ml4h.arguments import parse_args\n",
    "from ml4h.TensorMap import TensorMap, Interpretation\n",
    "from ml4h.tensor_generators import test_train_valid_tensor_generators\n",
    "from ml4h.recipes import compare_multimodal_scalar_task_models, train_multimodal_multitask\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Constants\n",
    "HD5_FOLDER = './mnist_tensors/'\n",
    "OUTPUT_FOLDER = './outputs/'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Python features we make lots of use of in this notebook:\n",
    "- F Strings\n",
    "- Callback Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TensorMaps\n",
    "The critical data structure in the ml4h codebase is the TensorMap.\n",
    "This abstraction provides a way to translate ***any*** kind of input data, into structured numeric tensors with clear semantics for interpretation and modeling.  TensorMaps guarantee a shape, a way to construct tensors of that shape from the HD5 files created during tensorization and a meaning to the values in the tensor that the TensorMap yields.\n",
    "\n",
    "For example, in the `mnist.py` file these TensorMaps are defined:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mnist_image_from_hd5(tm: TensorMap, hd5: h5py.File, dependents: Dict = {}) -> np.ndarray:\n",
    "    return np.array(hd5['mnist_image'])\n",
    "\n",
    "\n",
    "mnist_image = TensorMap('mnist_image', shape=(28, 28, 1), tensor_from_file=mnist_image_from_hd5)\n",
    "\n",
    "\n",
    "def mnist_label_from_hd5(tm: TensorMap, hd5: h5py.File, dependents: Dict = {}) -> np.ndarray:\n",
    "    one_hot = np.zeros(tm.shape, dtype=np.float32)\n",
    "    one_hot[int(hd5['mnist_label'][0])] = 1.0\n",
    "    return one_hot\n",
    "\n",
    "\n",
    "mnist_label = TensorMap(\n",
    "    'mnist_label', Interpretation.CATEGORICAL, tensor_from_file=mnist_label_from_hd5,\n",
    "    channel_map={f'digit_{i}': i for i in range(10)},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similiarly, in the `gatk.py` file we define tensors to encode data about genomic variants.  Specifically, we create 3 TensorMaps: `reference` is a 1-hot encoded 128 base-pair window of DNA sequence. `read_tensor` is an alignment of as many as 128 different DNA reads overlapping a 128 base-pair window of reference DNA.  This TensorMap includes 15 channels which encode the DNA bases from the reference from the read sequence and meta data belonging to each read.  Lastly, we define the `CATEGORICAL` TensorMap  `variant_label` which encodes the truth status of this particular genomic variant.  In this dataset we are considering on SNPs and small insertions or deletions giving us the 4 labels: `'NOT_SNP', 'NOT_INDEL', 'SNP', 'INDEL'`.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_data(dataset):\n",
    "    ''' Loads the dataset\n",
    "    :param dataset: the path to the dataset (here MNIST)'''\n",
    "    data_dir, data_file = os.path.split(dataset)\n",
    "    if data_dir == \"\" and not os.path.isfile(dataset):\n",
    "        # Check if dataset is in the data directory.\n",
    "        new_path = os.path.join(\"data\", dataset)\n",
    "        if os.path.isfile(new_path) or data_file == 'mnist.pkl.gz':\n",
    "            dataset = new_path\n",
    "\n",
    "    if (not os.path.isfile(dataset)) and data_file == 'mnist.pkl.gz':\n",
    "        from urllib.request import urlretrieve\n",
    "        origin = ('http://www.iro.umontreal.ca/~lisa/deep/data/mnist/mnist.pkl.gz')\n",
    "        print('Downloading data from %s' % origin)\n",
    "        if not os.path.exists(os.path.dirname(dataset)):\n",
    "            os.makedirs(os.path.dirname(dataset))\n",
    "        urlretrieve(origin, dataset)\n",
    "\n",
    "    print('loading data...')\n",
    "    f = gzip.open(dataset, 'rb')\n",
    "    if sys.version_info[0] == 3:\n",
    "        u = pickle._Unpickler(f)\n",
    "        u.encoding = 'latin1'\n",
    "        train_set, valid_set, test_set = u.load()\n",
    "    else:\n",
    "        train_set, valid_set, test_set = pickle.load(f)\n",
    "    f.close()\n",
    "\n",
    "    return train_set, valid_set, test_set\n",
    "\n",
    "\n",
    "def mnist_as_hd5(hd5_folder):\n",
    "    train, _, _ = load_data('mnist.pkl.gz')\n",
    "    mnist_images = train[0].reshape((-1, 28, 28, 1))\n",
    "    if not os.path.exists(hd5_folder):\n",
    "        os.makedirs(hd5_folder)\n",
    "    for i, mnist_image in enumerate(mnist_images):\n",
    "        with h5py.File(os.path.join(hd5_folder, f'{i}.hd5'), 'w') as hd5:\n",
    "            hd5.create_dataset('mnist_image', data=mnist_image)\n",
    "            hd5.create_dataset('mnist_label', data=[train[1][i]])\n",
    "        if (i+1) % 5000 == 0:\n",
    "            print(f'Wrote {i+1} MNIST images and labels as HD5 files')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the type of data used by the GATK tool CNNScoreVariants to filter DNA sequencing data.  The tensorization code is part of the GATK not ML4H, however tensorized data for use is available at: `gs://fc-500bd872-4a53-45c9-87d3-39656bd83f85/data/hg002_na24385_ml4h_tensors_v2021_10_14.tar.gz`. Once the data has been localized you can unpack the HD5 files into the `HD5_FOLDER` with the cell below (assuming the tar.gz file is in the same directory as:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnist_as_hd5(HD5_FOLDER)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Model Factory\n",
    "The function ***make_multimodal_multitask_model()*** takes lists of TensorMaps and connects them with intelligent goo.  Specifically, given a list of TensorMaps that are model inputs and TensorMaps that are desired outputs the model factory will build a model and loss appropriate for the dimensions and interpretations of the data at hand.  The depending on the input and output TensorMaps provided, the Model Factory will build models for many different situations including:\n",
    "- Classification\n",
    "- Regression\n",
    "- Multitask\n",
    "- Multimodal\n",
    "- Multimodal Multitask\n",
    "- Autoencoders\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.argv = ['train', \n",
    "            '--tensors', HD5_FOLDER, \n",
    "            '--input_tensors', 'mnist.mnist_image',\n",
    "            '--output_tensors', 'mnist.mnist_label',\n",
    "            '--batch_size', '16',\n",
    "            '--epochs', '12',\n",
    "            '--output_folder', OUTPUT_FOLDER,\n",
    "            '--id', 'learn_2d_cnn'\n",
    "           ]\n",
    "args = parse_args()\n",
    "metrics = train_multimodal_multitask(args)\n",
    "\n",
    "sys.argv = ['train',\n",
    "            '--tensors', HD5_FOLDER,\n",
    "            '--input_tensors', 'mnist.mnist_image',\n",
    "            '--output_tensors', 'mnist.mnist_label',\n",
    "            '--activation', 'mish',\n",
    "            '--dense_blocks', '64', '64', '64',\n",
    "            '--batch_size', '16',\n",
    "            '--epochs', '12',\n",
    "            '--output_folder', OUTPUT_FOLDER,\n",
    "            '--id', 'learn_2d_cnn2'\n",
    "           ]\n",
    "args = parse_args()\n",
    "metrics = train_multimodal_multitask(args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compare Models that have been trained for the same task (ie with the same output TensorMap)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.argv = ['compare_scalar', \n",
    "            '--tensors', HD5_FOLDER, \n",
    "            '--input_tensors', 'mnist.mnist_image',\n",
    "            '--output_tensors', 'mnist.mnist_label',\n",
    "            '--id', 'mnist_model_comparison',\n",
    "            '--output_folder', OUTPUT_FOLDER,\n",
    "            '--model_files', f'{OUTPUT_FOLDER}learn_2d_cnn/learn_2d_cnn.h5',\n",
    "                            f'{OUTPUT_FOLDER}learn_2d_cnn2/learn_2d_cnn2.h5',\n",
    "            '--test_steps', '100', \n",
    "            '--batch_size', '16',\n",
    "           ]\n",
    "args = parse_args()\n",
    "\n",
    "generate_train, generate_valid, generate_test = test_train_valid_tensor_generators(**args.__dict__)\n",
    "compare_multimodal_scalar_task_models(args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Custom Architectures\n",
    "The default architecture produced by the ModelFactory is based the [DenseNet](https://arxiv.org/abs/1608.06993) Convolutional Neural Network.  It is extremely customizable as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.argv = ['train', \n",
    "            '--tensors', HD5_FOLDER, \n",
    "            '--input_tensors', 'mnist.mnist_image',\n",
    "            '--output_tensors', 'mnit.mnist_label',\n",
    "            '--output_folder', OUTPUT_FOLDER,\n",
    "            '--activation', 'swish',\n",
    "            '--conv_layers', '32',\n",
    "            '--conv_width', '32', '32', '32',\n",
    "            '--dense_blocks', '32', '24', '16',\n",
    "            '--dense_layers', '32',  '32', \n",
    "            '--block_size', '4',\n",
    "            '--pool_x', '2',\n",
    "            '--pool_y', '2',\n",
    "            '--inspect_model',\n",
    "            '--epochs', '1',\n",
    "            '--batch_size', '4',\n",
    "            '--id', 'hypertuned_2d',\n",
    "           ]\n",
    "args = parse_args()\n",
    "generate_train, generate_valid, generate_test = test_train_valid_tensor_generators(**args.__dict__)\n",
    "train_multimodal_multitask(args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After running the cell above the diagram of the model architecture will be saved at: `./outputs/hypertuned_2d/architecture_graph_hypertuned_2d.png`"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}