Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Docker endpoint #10

Open
wants to merge 11 commits into
base: master
Choose a base branch
from
6 changes: 4 additions & 2 deletions .github/workflows/dockerbuild.yml
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ on:
- 'master'
pull_request:
branches:
- main
- master

jobs:
build:
Expand All @@ -22,6 +22,8 @@ jobs:

- name: Checkout repo
uses: actions/checkout@v2
with:
submodules: true

- name: Set up QEMU
uses: docker/setup-qemu-action@v1
Expand All @@ -37,6 +39,6 @@ jobs:
password: ${{secrets.CR_PAT}}

- name: Build and push docker image
run: sh ./Docker/dockerbuild.sh \
run: sh ./dockerbuild.sh \
"${{github.sha}}" \
"dipy"
1 change: 1 addition & 0 deletions docker/Dockerfile → Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ COPY run_dipy_cpu_hardi.py /opt/GPUStreamlines/run_dipy_cpu_hardi.py
COPY merge_trk.sh /opt/exec/merge_trk.sh
COPY cuslines /opt/GPUStreamlines/cuslines
COPY external /opt/GPUStreamlines/external

RUN mkdir -p /opt/exec/output

# compile
Expand Down
14 changes: 4 additions & 10 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -104,23 +104,17 @@ Destroy GPUTracker...

## Running on AWS with Docker
First, set up an AWS instance with GPU and ssh into it (we recommend a P3 instance with at least 1 V100 16 GB GPU and a Deep Learning AMI Ubuntu 18.04 v 33.0.). Then do the following:
1. Log in to GitHub docker registry:
1. Pull the container:
```
$ docker login -u <github id> docker.pkg.github.com
```
2. Enter your GitHub access token. If you do not have one, create it on the GitHub general security settings page and enable package read access for that token.
3. Pull the container:
```
$ docker pull docker.pkg.github.com/dipy/gpustreamlines/gpustreamlines:latest
$ docker pull docker pull ghcr.io/dipy/gpustreamlines:latest
```
4. Run the code, mounting the current directory into the container for easy result retrieval:
```
$ docker run --gpus=all -v ${PWD}:/opt/exec/output:rw -it docker.pkg.github.com/dipy/gpustreamlines/gpustreamlines:latest \
$ docker run --gpus=all -v ${PWD}:/opt/exec/output:rw -it ghcr.io/dipy/gpustreamlines:latest \
python run_dipy_gpu_hardi.py --chunk-size 100000 --ngpus 1 --output-prefix output/hardi_gpu_full --use-fast-write
```
5. The code produces a number of independent track files (one per processed "chunk"), but we have provided a merge script to combine them into a single trk file. To merge files, run:
```
$ docker run --gpus=all -v ${PWD}:/opt/exec/output:rw -it docker.pkg.github.com/dipy/gpustreamlines/gpustreamlines:latest \
$ docker run --gpus=all -v ${PWD}:/opt/exec/output:rw -it ghcr.io/dipy/gpustreamlines:latest \
./merge_trk.sh -o output/hardi_tracks.trk output/hardi_gpu_full*
```

4 changes: 2 additions & 2 deletions docker/dockerbuild.sh → dockerbuild.sh
Original file line number Diff line number Diff line change
@@ -1,13 +1,13 @@
COMMIT=${1}
COMMIT="$(echo "${COMMIT}" | tr -d '[:space:]')"
export COMMIT
NO_TAG="ghcr.io/${2}/GPUStreamlines"
NO_TAG="ghcr.io/${2}/gpustreamlines"
TAG="${NO_TAG}:${COMMIT}"
TAG2="${NO_TAG}:latest"
TAG="$(echo "${TAG}" | tr -d '[:space:]')"
TAG2="$(echo "${TAG2}" | tr -d '[:space:]')"
NO_TAG="$(echo "${NO_TAG}" | tr -d '[:space:]')"

echo $TAG
docker build --no-cache -t $TAG -t $TAG2 --build-arg COMMIT ./docker
docker build --no-cache -t $TAG -t $TAG2 --build-arg COMMIT .
docker push --all-tags $NO_TAG
123 changes: 62 additions & 61 deletions run_dipy_gpu.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,23 +31,19 @@
import argparse
import random
import time
from tqdm import trange, tqdm
from tqdm import tqdm

import numpy as np

import dipy.reconst.dti as dti
from dipy.io import read_bvals_bvecs
from dipy.io.stateful_tractogram import Origin, Space, StatefulTractogram
from dipy.io.image import load_nifti
from dipy.io.stateful_tractogram import Space, StatefulTractogram
from dipy.io.streamline import save_tractogram
from dipy.tracking import utils
from dipy.core.gradients import gradient_table
from dipy.data import small_sphere
#from dipy.direction import BootDirectionGetter
from dipy.reconst.shm import OpdtModel
#from dipy.tracking.local import LocalTracking, ThresholdTissueClassifier
from dipy.reconst import shm

import nibabel as nib
from nibabel.orientations import aff2axcodes

# Import custom module
Expand All @@ -59,39 +55,40 @@
np.random.seed(0)
random.seed(0)

#Get Gradient values
def get_gtab(fbval, fbvec):
bvals, bvecs = read_bvals_bvecs(fbval, fbvec)
gtab = gradient_table(bvals, bvecs)
return gtab

def get_img(ep2_seq):
img = nib.load(ep2_seq)
return img

print("parsing arguments")
parser = argparse.ArgumentParser()
parser.add_argument("nifti_file", help="path to the ep2multiband sequence nifti file")
parser.add_argument("nifti_file", help="path to the dwi nifti file")
parser.add_argument("bvals", help="path to the bvals")
parser.add_argument("bvecs", help="path to the bvecs")
parser.add_argument("mask_nifti", help="path to the mask file")
parser.add_argument("--fa_numpy", default=None, help="path to the FA numpy file")
parser.add_argument("--fa_file", default=None,
help="path to the FA file (nifti or numpy")
parser.add_argument("--roi_nifti", default=None, help="path to the ROI file")
parser.add_argument("--output-prefix", type=str, default ='', help="path to the output file")
parser.add_argument("--chunk-size", type=int, required=True, help="how many seeds to process per sweep, per GPU")
parser.add_argument("--sampling-density", type=int, default=3, help="sampling density for seed generation")
parser.add_argument("--ngpus", type=int, required=True, help="number of GPUs to use")
parser.add_argument("--use-fast-write", action='store_true', help="use fast file write")
parser.add_argument("--max-angle", type=float, default=1.0471975511965976, help="default: 60 deg (in rad)")
parser.add_argument("--min-signal", type=float, default=1.0, help="default: 1.0")
parser.add_argument("--tc-threshold", type=float, default=0.1, help="default: 0.1")
parser.add_argument("--step-size", type=float, default=0.5, help="default: 0.5")
parser.add_argument("--output-prefix", type=str, default='',
help="path to the output file")
parser.add_argument("--chunk-size", type=int, required=True,
help="how many seeds to process per sweep, per GPU")
parser.add_argument("--sampling-density", type=int, default=3,
help="sampling density for seed generation")
parser.add_argument("--ngpus", type=int, required=True,
help="number of GPUs to use")
parser.add_argument("--use-fast-write", action='store_true',
help="use fast file write")
parser.add_argument("--max-angle", type=float, default=1.0471975511965976,
help="default: 60 deg (in rad)")
parser.add_argument("--min-signal", type=float, default=1.0,
help="default: 1.0")
parser.add_argument("--tc-threshold", type=float, default=0.1,
help="default: 0.1")
parser.add_argument("--step-size", type=float, default=0.5,
help="default: 0.5")
args = parser.parse_args()

img = get_img(args.nifti_file)
img = load_nifti(args.nifti_file, return_img=True)
voxel_order = "".join(aff2axcodes(img.affine))
gtab = get_gtab(args.bvals, args.bvecs)
mask = get_img(args.mask_nifti)
gtab = gradient_table(args.bvals, args.bvecs)
mask = load_nifti(args.mask_nifti, return_img=True)
data = img.get_fdata()

# resample mask if necessary
Expand All @@ -102,7 +99,6 @@ def get_img(ep2_seq):
img.shape[:3], img.affine,
mask.shape[:3], mask.affine)
mask = affine_map.transform(mask.get_fdata())
#mask = np.round(mask)
else:
mask = mask.get_fdata()

Expand All @@ -122,28 +118,28 @@ def get_img(ep2_seq):
np.save(args.fa_numpy, FA)

# Setup tissue_classifier args
#tissue_classifier = ThresholdTissueClassifier(FA, 0.1)
metric_map = np.asarray(FA, 'float64')

# resample roi if necessary
if args.roi_nifti is not None:
roi = get_img(args.roi_nifti)
roi_data = roi.get_fdata()
roi_data, roi = load_nifti(args.roi_nifti,
return_img=True,
as_ndarray=True)
else:
roi_data = np.zeros(data.shape[:3])
roi_data[ FA > 0.1 ] = 1.
roi_data[FA > 0.1] = 1.

# Create seeds for ROI
seed_mask = utils.seeds_from_mask(roi_data, density=args.sampling_density, affine=np.eye(4))
seed_mask = utils.seeds_from_mask(roi_data,
density=args.sampling_density,
affine=np.eye(4))

# Setup model
print('slowadcodf')
print('OpdtModel')
sh_order = 6
model = OpdtModel(gtab, sh_order=sh_order, min_signal=1)

# Setup direction getter args
print('Bootstrap direction getter')
#boot_dg = BootDirectionGetter.from_data(data, model, max_angle=60., sphere=small_sphere)
b0s_mask = gtab.b0s_mask
dwi_mask = ~b0s_mask

Expand All @@ -157,7 +153,7 @@ def get_img(ep2_seq):
phi = sphere.phi
sampling_matrix, _, _ = shm.real_sym_sh_basis(sh_order, theta, phi)

## from BootPmfGen __init__
# from BootPmfGen __init__
# setup H and R matrices
# TODO: figure out how to get H, R matrices from direction getter object
x, y, z = model.gtab.gradients[dwi_mask].T
Expand All @@ -173,15 +169,15 @@ def get_img(ep2_seq):
global_chunk_size = args.chunk_size * args.ngpus
nchunks = (seed_mask.shape[0] + global_chunk_size - 1) // global_chunk_size

#streamline_generator = LocalTracking(boot_dg, tissue_classifier, seed_mask, affine=np.eye(4), step_size=.5)
gpu_tracker = cuslines.GPUTracker(args.max_angle,
args.min_signal,
args.tc_threshold,
args.step_size,
dataf, H, R, delta_b, delta_q,
b0s_mask.astype(np.int32), metric_map, sampling_matrix,
sphere.vertices, sphere.edges.astype(np.int32),
ngpus=args.ngpus, rng_seed=0)
gpu_tracker = cuslines.GPUTracker(
args.max_angle,
args.min_signal,
args.tc_threshold,
args.step_size,
dataf, H, R, delta_b, delta_q,
b0s_mask.astype(np.int32), metric_map, sampling_matrix,
sphere.vertices, sphere.edges.astype(np.int32),
ngpus=args.ngpus, rng_seed=0)
t1 = time.time()
streamline_time = 0
io_time = 0
Expand All @@ -191,30 +187,35 @@ def get_img(ep2_seq):
for idx in range(int(nchunks)):
# Main streamline computation
ts = time.time()
streamlines = gpu_tracker.generate_streamlines(seed_mask[idx*global_chunk_size:(idx+1)*global_chunk_size])
streamlines = gpu_tracker.generate_streamlines(
seed_mask[idx * global_chunk_size:(idx+1) * global_chunk_size])
te = time.time()
streamline_time += (te-ts)
pbar.update(seed_mask[idx*global_chunk_size:(idx+1)*global_chunk_size].shape[0])
#print("Generated {} streamlines from {} seeds, time: {} s".format(len(streamlines),
# seed_mask[idx*global_chunk_size:(idx+1)*global_chunk_size].shape[0],
# te-ts))
pbar.update(
seed_mask[
idx * global_chunk_size:(idx+1) * global_chunk_size].shape[0])

# Save tracklines file
if args.output_prefix:
if args.use_fast_write:
prefix = "{}.{}_{}".format(args.output_prefix, idx+1, nchunks)
ts = time.time()
gpu_tracker.dump_streamlines(prefix, voxel_order, roi_data.shape, img.header.get_zooms(), img.affine)
gpu_tracker.dump_streamlines(
prefix, voxel_order, roi_data.shape,
img.header.get_zooms(), img.affine)
te = time.time()
print("Saved streamlines to {}_*.trk, time {} s".format(prefix, te-ts))
print("Saved streamlines to {}_*.trk, time {} s".format(
prefix, te-ts))
else:
fname = "{}.{}_{}.trk".format(args.output_prefix, idx+1, nchunks)
fname = "{}.{}_{}.trk".format(
args.output_prefix, idx+1, nchunks)
ts = time.time()
#save_tractogram(fname, streamlines, img.affine, vox_size=roi.header.get_zooms(), shape=roi_data.shape)
sft = StatefulTractogram(streamlines, args.nifti_file, Space.VOX)
sft = StatefulTractogram(
streamlines, args.nifti_file, Space.VOX)
save_tractogram(sft, fname)
te = time.time()
print("Saved streamlines to {}, time {} s".format(fname, te-ts))
print("Saved streamlines to {}, time {} s".format(
fname, te-ts))
io_time += (te-ts)

pbar.close()
Expand All @@ -225,4 +226,4 @@ def get_img(ep2_seq):
print("Streamline generation total time: {} sec".format(t2-t1))
print("\tStreamline processing: {} sec".format(streamline_time))
if args.output_prefix:
print("\tFile writing: {} sec".format(io_time))
print("\tFile writing: {} sec".format(io_time))