id_gpu

Author: qsyao

cuda_backend/kernel/test.cu

@@ -2,14 +2,14 @@
 
 extern "C" {
 
-Handle* init_2D_handle(size_t y, size_t x, int mode_type, float c_val){
-    Handle *ret = new Handle(mode_type, c_val);
+Handle* init_2D_handle(size_t y, size_t x, int mode_type, float c_val, int id_gpu){
+    Handle *ret = new Handle(mode_type, c_val, id_gpu);
     ret->set_2D(y, x);
     return ret;
 }
 
-Handle* init_3D_handle(size_t z, size_t y, size_t x, int mode_type, float c_val){
-    Handle *ret = new Handle(mode_type, c_val);
+Handle* init_3D_handle(size_t z, size_t y, size_t x, int mode_type, float c_val, int id_gpu){
+    Handle *ret = new Handle(mode_type, c_val, id_gpu);
     ret->set_3D(z, y, x);
     return ret;
 }

cuda_backend/kernel/utils.cuh

@@ -18,8 +18,9 @@ extern "C"{
 
 class Handle {
 public:
-    Handle(int mode_type, float c_val) : batchsize(1), dim_x(1), dim_y(1),
+    Handle(int mode_type, float c_val, int id_gpu) : batchsize(1), dim_x(1), dim_y(1),
                  dim_z(1), mode_type(mode_type), c_val(c_val){
+        checkCudaErrors(cudaSetDevice(id_gpu));
         checkCudaErrors(cudaStreamCreate(&stream));
         checkCudaErrors(curandCreateGenerator(&gen, 
                        CURAND_RNG_PSEUDO_DEFAULT));

cuda_backend/py_api.py

@@ -7,11 +7,11 @@
 lib = CDLL(lib_dir + "/libcudaAugmentation.so", RTLD_GLOBAL)
 
 init_2D = lib.init_2D_handle
-init_2D.argtypes = [c_int, c_int, c_int, c_float]
+init_2D.argtypes = [c_int, c_int, c_int, c_float, c_int]
 init_2D.restype = c_void_p
 
 init_3D = lib.init_3D_handle
-init_3D.argtypes = [c_int, c_int, c_int, c_int, c_float]
+init_3D.argtypes = [c_int, c_int, c_int, c_int, c_float, c_int]
 init_3D.restype = c_void_p
 
 l_i = lib.linear_interpolate
@@ -167,7 +167,7 @@ def defrom(self, handle):
         return None
 
 class Handle(object):
-    def __init__(self, shape, RGB=False, mode='constant', cval=0.0):
+    def __init__(self, shape, RGB=False, mode='constant', cval=0.0, id_gpu=0):
         self.RGB = RGB
         self.shape = shape
         if self.RGB:
@@ -193,11 +193,11 @@ def __init__(self, shape, RGB=False, mode='constant', cval=0.0):
             raise ValueError
 
         if(len(shape) == 2 or RGB):
-            self.cuda_handle = init_2D(self.shape[0], self.shape[1], type_mode, float(cval))
+            self.cuda_handle = init_2D(self.shape[0], self.shape[1], type_mode, float(cval), id_gpu)
             self.is_3D = False
         else:
             self.is_3D = True
-            self.cuda_handle = init_3D(shape[0], shape[1],  shape[2], type_mode, float(cval))
+            self.cuda_handle = init_3D(shape[0], shape[1],  shape[2], type_mode, float(cval), id_gpu)
 
     def augment(self, img):
         if self.RGB:

deform.py

@@ -1,3 +1,6 @@
+'''
+Reference to https://github.com/MIC-DKFZ/batchgenerators
+'''
 
 import numpy as np
 import SimpleITK as sitk
@@ -68,8 +71,8 @@ def spatial_augment(img, RGB=False, do_scale=True, scale=0.5, angle=0.75*np.pi,
     coords = create_zero_centered_coordinate_mesh(img.shape)
 
     # coords = scale_coords(coords, scale)
-    # coords = rotate_coords_3d(coords, angle, angle, angle)
-    coords = elastic_deform_coordinates(coords, 500, 12)
+    coords = rotate_coords_3d(coords, angle, angle, angle)
+    # coords = elastic_deform_coordinates(coords, 500, 12)
 
     for d in range(len(img.shape)):
         ctr = float(np.round(img.shape[d] / 2.))

python_augmentation.py

@@ -12,13 +12,6 @@
 
 np.set_printoptions(precision=3)
 
-def create_zero_centered_coordinate_mesh(shape):
-    tmp = tuple([np.arange(i) for i in shape])
-    coords = np.array(np.meshgrid(*tmp, indexing='ij')).astype(float)
-    for d in range(len(shape)):
-        coords[d] -= ((np.array(shape).astype(float)) / 2.)[d]
-    return coords
-
 def check(correct, output):
     '''
     Unit Test Pass When less than 0.01 rate pixels loss ( > 0.001)