def backward(ctx, grad_out): astra.set_gpu_index([grad_out.device.index], memory=10 * 1024 * 1024 * 1024) grad_input = gradProjSize = gradVectors = None if ctx.needs_input_grad[0]: # permute here to handle opposite data layouts bproj_id, bproj_data = astra.create_backprojection3d_gpu( grad_out.squeeze(1).permute(2, 0, 1).data.cpu().numpy(), ctx.proj_geom, ctx.vol_geom) astra.data3d.delete(bproj_id) grad_input = Tensor(bproj_data).cuda(non_blocking=True).permute( 2, 1, 0) return grad_input, gradProjSize, gradVectors
def import_astra_GPU(): import astra sp = subprocess.Popen(['nvidia-smi', '-q'], stdout=subprocess.PIPE, stderr=subprocess.PIPE) out_str = sp.communicate() out_list = out_str[0].decode("utf-8").split('\n') out_dict = {} for item in out_list: try: key, val = item.split(':') key, val = key.strip(), val.strip() out_dict[key] = val except: pass num_gpu = int(out_dict['Attached GPUs']) astra.set_gpu_index([i for i in range(num_gpu)])
""" import numpy as np import ddf_fdk as ddf #ddf.import_astra_GPU() import nn_fdk as nn from sacred import Experiment from sacred.observers import FileStorageObserver import gc import pylab import os import time import h5py import sys sys.path.append('../nn_fdk/') import astra astra.set_gpu_index([0, 1, 2, 3]) import load_and_preprocess_CA as cap # %% path = 'python_data/results/' ex = Experiment() # %% @ex.config def cfg(): it_i = 1 it_j = 21 bp = '/bigstore/lagerwer/data/FleXray/' path = f'{bp}Walnuts/Walnut{it_j}/Projections/' dset = f'tubeV{2}' pd = 'processed_data/'
# The ASTRA Toolbox is distributed in the hope that it will be useful, # but WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # GNU General Public License for more details. # # You should have received a copy of the GNU General Public License # along with the ASTRA Toolbox. If not, see <http://www.gnu.org/licenses/>. # # ----------------------------------------------------------------------- import astra import numpy as np # Set up multi-GPU usage. # This only works for 3D GPU forward projection and back projection. astra.set_gpu_index([0, 1]) # Optionally, you can also restrict the amount of GPU memory ASTRA will use. # The line commented below sets this to 1GB. #astra.set_gpu_index([0,1], memory=1024*1024*1024) vol_geom = astra.create_vol_geom(1024, 1024, 1024) angles = np.linspace(0, np.pi, 1024, False) proj_geom = astra.create_proj_geom('parallel3d', 1.0, 1.0, 1024, 1024, angles) # Create a simple hollow cube phantom cube = np.zeros((1024, 1024, 1024)) cube[128:895, 128:895, 128:895] = 1 cube[256:767, 256:767, 256:767] = 0
def forward(ctx, input, projSize, vectors): #angles=T(B,3) """ input - GPU tensor ctx.vectors - numProj x 12 following astra convention ctx.projSize - size in pixels of projections returns: proj: GPU tensor numProj x 1 x projSize x projSize avoid passing back to the cpu based on https://github.com/astra-toolbox/astra-toolbox/blob/master/samples/python/s021_pygpu.py and https://github.com/astra-toolbox/astra-toolbox/blob/10d87f45bc9311c0408e4cacdec587eff8bc37f8/python/astra/creators.py create_sino3d_gpu """ # pytorch is row major (later dimensions change faster), while astra is col major (opposite) # to counteract this, we need to transpose # the contiguous is needed to change the memory layout astra.set_gpu_index([input.device.index], memory=1 * 1024 * 1024 * 1024) input = input.permute(2, 1, 0).contiguous() # setup volume x, y, z = input.shape strideBytes = input.stride( -2) * 32 / 8 # 32 bits in a float, 8 bits in a byte vol_link = astra.data3d.GPULink(input.data_ptr(), x, y, z, strideBytes) vol_geom = astra.create_vol_geom(x, y, z) vol_id = astra.data3d.link('-vol', vol_geom, vol_link) # setup projection proj = torch.empty(projSize, len(vectors), projSize, dtype=torch.float, device='cuda') x, y, z = proj.shape strideBytes = proj.stride( -2) * 32 / 8 # 32 bits in a float, 8 bits in a byte proj_link = astra.data3d.GPULink(proj.data_ptr(), x, y, z, strideBytes) proj_geom = astra.create_proj_geom('parallel3d_vec', projSize, projSize, vectors.numpy()) proj_id = astra.data3d.link('-sino', proj_geom, proj_link) cfg = astra.astra_dict('FP3D_CUDA') cfg['VolumeDataId'] = vol_id cfg['ProjectionDataId'] = proj_id alg_id = astra.algorithm.create(cfg) astra.algorithm.run(alg_id) # again handle the row vs col major problem # the output is y_0 x numProj x y_1, which pytorch interprets in the reverse order proj = proj.permute(1, 2, 0).unsqueeze(1) # store data for back propagation ctx.proj_geom = proj_geom ctx.vol_geom = vol_geom return proj
#it under the terms of the GNU General Public License as published by #the Free Software Foundation, either version 3 of the License, or #(at your option) any later version. # #The Python interface to the ASTRA Toolbox is distributed in the hope that it will be useful, #but WITHOUT ANY WARRANTY; without even the implied warranty of #MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the #GNU General Public License for more details. # #You should have received a copy of the GNU General Public License #along with the Python interface to the ASTRA Toolbox. If not, see <http://www.gnu.org/licenses/>. # #----------------------------------------------------------------------- import matlab as m from creators import astra_dict,create_vol_geom, create_proj_geom, create_backprojection, create_sino, create_reconstruction, create_projector,create_sino3d_gpu, create_backprojection3d_gpu from functions import data_op, add_noise_to_sino,clear from extrautils import clipCircle from ASTRAProjector import ASTRAProjector2D import data2d import astra import data3d import algorithm import projector import matrix import os try: astra.set_gpu_index(int(os.environ['ASTRA_GPU_INDEX'])) except KeyError: pass
def fwd_gpu_recon(func, i, num_gpus, cors, **kwargs): import astra astra.set_gpu_index(i % num_gpus) second_shared_data[i] = func(shared_data[i], cors[i], **kwargs)
# The ASTRA Toolbox is distributed in the hope that it will be useful, # but WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # GNU General Public License for more details. # # You should have received a copy of the GNU General Public License # along with the ASTRA Toolbox. If not, see <http://www.gnu.org/licenses/>. # # ----------------------------------------------------------------------- import astra import numpy as np # Set up multi-GPU usage. # This only works for 3D GPU forward projection and back projection. astra.set_gpu_index([0,1]) # Optionally, you can also restrict the amount of GPU memory ASTRA will use. # The line commented below sets this to 1GB. #astra.set_gpu_index([0,1], memory=1024*1024*1024) vol_geom = astra.create_vol_geom(1024, 1024, 1024) angles = np.linspace(0, np.pi, 1024,False) proj_geom = astra.create_proj_geom('parallel3d', 1.0, 1.0, 1024, 1024, angles) # Create a simple hollow cube phantom cube = np.zeros((1024,1024,1024)) cube[128:895,128:895,128:895] = 1 cube[256:767,256:767,256:767] = 0