Skip to content

Commit

Permalink
112 combining spatial transforms (#131)
Browse files Browse the repository at this point in the history 
* affine transforms
* randomised affine and elastic deformation
* add unit tests
* add a 2D notebook demo
* add speed demo
  • Loading branch information
@wyli
wyli authored Mar 9, 2020
1 parent 1e230ed commit 7304182
Show file tree
Hide file tree
Showing 16 changed files with 2,164 additions and 2 deletions.
372 changes: 372 additions & 0 deletions examples/transform_speed.ipynb
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,372 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Data loading pipeline examples\n",
"\n",
"The purpose of this notebook is to illustrate reading Nifti files and test speed of different methods."
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"MONAI version: 0.0.1\n",
"Python version: 3.5.6 |Anaconda, Inc.| (default, Aug 26 2018, 16:30:03) [GCC 4.2.1 Compatible Clang 4.0.1 (tags/RELEASE_401/final)]\n",
"Numpy version: 1.18.1\n",
"Pytorch version: 1.4.0\n",
"Ignite version: 0.3.0\n"
]
}
],
"source": [
"%matplotlib inline\n",
"\n",
"import os\n",
"import sys\n",
"from glob import glob\n",
"import tempfile\n",
"\n",
"import numpy as np\n",
"import nibabel as nib\n",
"\n",
"\n",
"import torch\n",
"from torch.utils.data import DataLoader\n",
"from torch.multiprocessing import Pool, Process, set_start_method\n",
"try:\n",
" set_start_method('spawn')\n",
"except RuntimeError:\n",
" pass\n",
"\n",
"sys.path.append('..') # assumes this is where MONAI is\n",
"\n",
"import monai\n",
"from monai.transforms.compose import Compose\n",
"from monai.data.nifti_reader import NiftiDataset\n",
"from monai.transforms import (AddChannel, Rescale, ToTensor, \n",
" UniformRandomPatch, Rotate, RandAffine)\n",
"\n",
"monai.config.print_config()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 0. Preparing input data (nifti images)\n",
"\n",
"Create a number of test Nifti files, 3d single channel images with spatial size (256, 256, 256) voxels."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"tempdir = tempfile.mkdtemp()\n",
"\n",
"for i in range(5):\n",
" im, seg = monai.data.synthetic.create_test_image_3d(256,256,256)\n",
" \n",
" n = nib.Nifti1Image(im, np.eye(4))\n",
" nib.save(n, os.path.join(tempdir, 'im%i.nii.gz'%i))\n",
" \n",
" n = nib.Nifti1Image(seg, np.eye(4))\n",
" nib.save(n, os.path.join(tempdir, 'seg%i.nii.gz'%i))"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"# prepare list of image names and segmentation names\n",
"images = sorted(glob(os.path.join(tempdir,'im*.nii.gz')))\n",
"segs = sorted(glob(os.path.join(tempdir,'seg*.nii.gz')))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 1. Test image loading with minimal preprocessing"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"torch.Size([3, 1, 256, 256, 256]) torch.Size([3, 1, 256, 256, 256])\n"
]
}
],
"source": [
"imtrans = Compose([\n",
" AddChannel(),\n",
" ToTensor()\n",
"]) \n",
"\n",
"segtrans = Compose([\n",
" AddChannel(),\n",
" ToTensor()\n",
"]) \n",
" \n",
"ds = NiftiDataset(images, segs, transform=imtrans, seg_transform=segtrans)\n",
"loader = DataLoader(ds, batch_size=3, num_workers=8)\n",
"\n",
"im, seg = monai.utils.misc.first(loader)\n",
"print(im.shape, seg.shape)"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"5.11 s ± 207 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
]
}
],
"source": [
"%timeit data = next(iter(loader))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 2. Test image-patch loading with CPU multi-processing:\n",
"\n",
"- rotate (256, 256, 256)-voxel in the plane axes=(1, 2)\n",
"- extract random (64, 64, 64) patches\n",
"- implemented in MONAI using ` scipy.ndimage.rotate`"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"torch.Size([3, 1, 64, 64, 64]) torch.Size([3, 1, 64, 64, 64])\n"
]
}
],
"source": [
"images = sorted(glob(os.path.join(tempdir,'im*.nii.gz')))\n",
"segs = sorted(glob(os.path.join(tempdir,'seg*.nii.gz')))\n",
"\n",
"imtrans = Compose([\n",
" Rescale(),\n",
" AddChannel(),\n",
" Rotate(angle=45.),\n",
" UniformRandomPatch((64, 64, 64)),\n",
" ToTensor()\n",
"]) \n",
"\n",
"segtrans = Compose([\n",
" AddChannel(),\n",
" Rotate(angle=45.),\n",
" UniformRandomPatch((64, 64, 64)),\n",
" ToTensor()\n",
"]) \n",
" \n",
"ds = NiftiDataset(images, segs, transform=imtrans, seg_transform=segtrans)\n",
"loader = DataLoader(ds, batch_size=3, num_workers=8, pin_memory=torch.cuda.is_available())\n",
"\n",
"im, seg = monai.utils.misc.first(loader)\n",
"print(im.shape, seg.shape)"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"10.3 s ± 175 ms per loop (mean ± std. dev. of 7 runs, 3 loops each)\n"
]
}
],
"source": [
"%timeit -n 3 data = next(iter(loader))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"(the above results were based on a 2.9 GHz 6-Core Intel Core i9)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 3. Test image-patch loading with preprocessing on GPU:\n",
"\n",
"- random rotate (256, 256, 256)-voxel in the plane axes=(1, 2)\n",
"- extract random (64, 64, 64) patches\n",
"- implemented in MONAI using native pytorch resampling"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"torch.Size([3, 1, 64, 64, 64]) torch.Size([3, 1, 64, 64, 64])\n"
]
}
],
"source": [
"images = sorted(glob(os.path.join(tempdir,'im*.nii.gz')))\n",
"segs = sorted(glob(os.path.join(tempdir,'seg*.nii.gz')))\n",
"\n",
"# same parameter with different interpolation mode for image and segmentation\n",
"rand_affine_img = RandAffine(prob=1.0, rotate_range=np.pi/4, translate_range=(96, 96, 96),\n",
" spatial_size=(64, 64, 64), mode='bilinear',\n",
" as_tensor_output=True, device=torch.device('cuda:0'))\n",
"rand_affine_seg = RandAffine(prob=1.0, rotate_range=np.pi/4, translate_range=(96, 96, 96),\n",
" spatial_size=(64, 64, 64), mode='nearest',\n",
" as_tensor_output=True, device=torch.device('cuda:0'))\n",
" \n",
"imtrans = Compose([\n",
" Rescale(),\n",
" AddChannel(),\n",
" rand_affine_img,\n",
"]) \n",
"\n",
"segtrans = Compose([\n",
" AddChannel(),\n",
" rand_affine_seg,\n",
"]) \n",
" \n",
"ds = NiftiDataset(images, segs, transform=imtrans, seg_transform=segtrans)\n",
"loader = DataLoader(ds, batch_size=3, num_workers=0)\n",
"\n",
"im, seg = monai.utils.misc.first(loader)\n",
"\n",
"print(im.shape, seg.shape)"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"1.42 s ± 1.72 ms per loop (mean ± std. dev. of 7 runs, 3 loops each)\n"
]
}
],
"source": [
"%timeit -n 3 data = next(iter(loader))"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"TITAN Xp COLLECTORS EDITION\n",
"|===========================================================================|\n",
"| PyTorch CUDA memory summary, device ID 0 |\n",
"|---------------------------------------------------------------------------|\n",
"| CUDA OOMs: 0 | cudaMalloc retries: 0 |\n",
"|===========================================================================|\n",
"| Metric | Cur Usage | Peak Usage | Tot Alloc | Tot Freed |\n",
"|---------------------------------------------------------------------------|\n",
"| Allocated memory | 6144 KB | 156672 KB | 16680 MB | 16674 MB |\n",
"|---------------------------------------------------------------------------|\n",
"| Active memory | 6144 KB | 156672 KB | 16680 MB | 16674 MB |\n",
"|---------------------------------------------------------------------------|\n",
"| GPU reserved memory | 225280 KB | 225280 KB | 225280 KB | 0 B |\n",
"|---------------------------------------------------------------------------|\n",
"| Non-releasable memory | 14336 KB | 77824 KB | 11219 MB | 11205 MB |\n",
"|---------------------------------------------------------------------------|\n",
"| Allocations | 2 | 14 | 2222 | 2220 |\n",
"|---------------------------------------------------------------------------|\n",
"| Active allocs | 2 | 14 | 2222 | 2220 |\n",
"|---------------------------------------------------------------------------|\n",
"| GPU reserved segments | 8 | 8 | 8 | 0 |\n",
"|---------------------------------------------------------------------------|\n",
"| Non-releasable allocs | 1 | 6 | 1460 | 1459 |\n",
"|===========================================================================|\n",
"\n"
]
}
],
"source": [
"print(torch.cuda.get_device_name(0))\n",
"print(torch.cuda.memory_summary(0, abbreviated=True))"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [],
"source": [
"!rm -rf {tempdir}"
]
}
],
"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.5.6"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
Loading

0 comments on commit 7304182

Please sign in to comment.