def bench_fanbeam_forward(phantom, det_count, num_angles, source_dist, det_dist, warmup, repeats): angles = np.linspace(0, np.pi, num_angles, endpoint=False) radon = RadonFanbeam(phantom.size(1), angles, source_dist, det_dist, det_count) f = lambda x: radon.forward(x) return benchmark(f, phantom, warmup, repeats)
def test_fanbeam_error(device, batch_size, image_size, angles, spacing, distances, det_count, clip_to_circle): # generate random images # generate random images det_count = int(det_count * image_size) mask_radius = det_count / 2.0 if clip_to_circle else -1 x = generate_random_images(1, image_size, mask_radius)[0] s_dist, d_dist = distances s_dist *= image_size d_dist *= image_size # astra vol_geom = astra.create_vol_geom(x.shape[0], x.shape[1]) proj_geom = astra.create_proj_geom('fanflat', spacing, det_count, angles, s_dist, d_dist) proj_id = astra.create_projector('cuda', proj_geom, vol_geom) id, astra_y = astra.create_sino(x, proj_id) _, astra_bp = astra.create_backprojection(astra_y, proj_id) if clip_to_circle: astra_bp *= circle_mask(image_size, mask_radius) # TODO clean astra structures # our implementation radon = RadonFanbeam(image_size, angles, s_dist, d_dist, det_count=det_count, det_spacing=spacing, clip_to_circle=clip_to_circle) x = torch.FloatTensor(x).to(device).view(1, x.shape[0], x.shape[1]) # repeat data to fill batch size x = torch.cat([x] * batch_size, dim=0) our_fp = radon.forward(x) our_bp = radon.backprojection(our_fp) forward_error = relative_error(astra_y, our_fp[0].cpu().numpy()) back_error = relative_error(astra_bp, our_bp[0].cpu().numpy()) # if back_error > 5e-3: # plt.imshow(astra_bp) # plt.figure() # plt.imshow(our_bp[0].cpu().numpy()) # plt.show() print(np.max(our_fp.cpu().numpy()), np.max(our_bp.cpu().numpy())) print( f"batch: {batch_size}, size: {image_size}, angles: {len(angles)}, spacing: {spacing}, distances: {distances} circle: {clip_to_circle}, forward: {forward_error}, back: {back_error}" ) # TODO better checks assert_less(forward_error, 1e-2) assert_less(back_error, 5e-3)
def init_radon(self, beam, circle, det_dist): if beam == 'parallel': angles = np.linspace(0, np.pi, self.n_angles, endpoint=False) self.radon = Radon(self.img_size, angles, clip_to_circle=circle) self.radon_sparse = Radon(self.img_size, angles[::self.sample_ratio], clip_to_circle=circle) elif beam == 'fan': angles = np.linspace(0, self.n_angles / 180 * np.pi, self.n_angles, False) self.radon = RadonFanbeam(self.img_size, angles, source_distance=det_dist[0], det_distance=det_dist[1], clip_to_circle=circle, det_count=self.det_size) self.radon_sparse = RadonFanbeam(self.img_size, angles[::self.sample_ratio], source_distance=det_dist[0], det_distance=det_dist[1], clip_to_circle=circle, det_count=self.det_size) else: raise Exception('projection beam type undefined!') self.n_angles_sparse = len(angles[::self.sample_ratio])
def test_half(device, batch_size, image_size, angles, spacing, distances, det_count, clip_to_circle): # generate random images det_count = int(det_count * image_size) mask_radius = det_count / 2.0 if clip_to_circle else -1 x = generate_random_images(batch_size, image_size, mask_radius) s_dist, d_dist = distances s_dist *= image_size d_dist *= image_size # our implementation radon = RadonFanbeam(image_size, angles, s_dist, d_dist, det_count=det_count, det_spacing=spacing, clip_to_circle=clip_to_circle) x = torch.FloatTensor(x).to(device) # divide by len(angles) to avoid half-precision overflow sinogram = radon.forward(x) / len(angles) single_precision = radon.backprojection(sinogram) h_sino = radon.forward(x.half()) / len(angles) half_precision = radon.backprojection(h_sino) print(torch.min(half_precision).item(), torch.max(half_precision).item()) forward_error = relative_error(sinogram.cpu().numpy(), h_sino.cpu().numpy()) back_error = relative_error(single_precision.cpu().numpy(), half_precision.cpu().numpy()) print( f"batch: {batch_size}, size: {image_size}, angles: {len(angles)}, spacing: {spacing}, circle: {clip_to_circle}, forward: {forward_error}, back: {back_error}") assert_less(forward_error, 1e-3) assert_less(back_error, 1e-3)
def main(): parser = argparse.ArgumentParser( description='Benchmark and compare with Astra Toolbox') parser.add_argument('--task', default="all") parser.add_argument('--image-size', default=256, type=int) parser.add_argument('--angles', default=-1, type=int) parser.add_argument('--batch-size', default=32, type=int) parser.add_argument('--samples', default=50, type=int) parser.add_argument('--warmup', default=10, type=int) parser.add_argument('--output', default="") parser.add_argument('--circle', action='store_true') args = parser.parse_args() if args.angles == -1: args.angles = args.image_size device = torch.device("cuda") angles = np.linspace(0, 2 * np.pi, args.angles, endpoint=False).astype(np.float32) radon = Radon(args.image_size, angles, clip_to_circle=args.circle) radon_fb = RadonFanbeam(args.image_size, angles, args.image_size, clip_to_circle=args.circle) astra_pw = AstraParallelWrapper(angles, args.image_size) astra_fw = AstraFanbeamWrapper(angles, args.image_size) # astra = AstraWrapper(angles) if args.task == "all": tasks = ["forward", "backward", "fanbeam forward", "fanbeam backward"] elif args.task == "shearlet": # tasks = ["shearlet forward", "shearlet backward"] benchmark_shearlet(args) return else: tasks = [args.task] astra_fps = [] radon_fps = [] radon_half_fps = [] if "forward" in tasks: print("Benchmarking forward from device") x = generate_random_images(args.batch_size, args.image_size) dx = torch.FloatTensor(x).to(device) astra_time = benchmark_function(lambda y: astra_pw.forward(y), dx, args.samples, args.warmup) radon_time = benchmark_function(lambda y: radon.forward(y), dx, args.samples, args.warmup, sync=True) radon_half_time = benchmark_function(lambda y: radon.forward(y), dx.half(), args.samples, args.warmup, sync=True) astra_fps.append(args.batch_size / astra_time) radon_fps.append(args.batch_size / radon_time) radon_half_fps.append(args.batch_size / radon_half_time) print("Speedup:", astra_time / radon_time) print("Speedup half-precision:", astra_time / radon_half_time) print() if "backward" in tasks: print("Benchmarking backward from device") x = generate_random_images(args.batch_size, args.image_size) dx = torch.FloatTensor(x).to(device) astra_time = benchmark_function(lambda y: astra_pw.backward(y), dx, args.samples, args.warmup) radon_time = benchmark_function(lambda y: radon.backward(y), dx, args.samples, args.warmup, sync=True) radon_half_time = benchmark_function(lambda y: radon.backward(y), dx.half(), args.samples, args.warmup, sync=True) astra_fps.append(args.batch_size / astra_time) radon_fps.append(args.batch_size / radon_time) radon_half_fps.append(args.batch_size / radon_half_time) print("Speedup:", astra_time / radon_time) print("Speedup half-precision:", astra_time / radon_half_time) print() if "fanbeam forward" in tasks: print("Benchmarking fanbeam forward") x = generate_random_images(args.batch_size, args.image_size) dx = torch.FloatTensor(x).to(device) # astra_time = benchmark_function(lambda y: astra_fw.forward(y), dx, args.samples, args.warmup) radon_time = benchmark_function(lambda y: radon_fb.forward(y), dx, args.samples, args.warmup, sync=True) radon_half_time = benchmark_function(lambda y: radon_fb.forward(y), dx.half(), args.samples, args.warmup, sync=True) astra_fps.append(args.batch_size / astra_time) radon_fps.append(args.batch_size / radon_time) radon_half_fps.append(args.batch_size / radon_half_time) print("Speedup:", astra_time / radon_time) print("Speedup half-precision:", astra_time / radon_half_time) print() if "fanbeam backward" in tasks: print("Benchmarking fanbeam backward") x = generate_random_images(args.batch_size, args.image_size) dx = torch.FloatTensor(x).to(device) # astra_time = benchmark_function(lambda y: astra_fw.backward(y), dx, args.samples, args.warmup) radon_time = benchmark_function(lambda y: radon_fb.backprojection(y), dx, args.samples, args.warmup, sync=True) radon_half_time = benchmark_function( lambda y: radon_fb.backprojection(y), dx.half(), args.samples, args.warmup, sync=True) astra_fps.append(args.batch_size / astra_time) radon_fps.append(args.batch_size / radon_time) radon_half_fps.append(args.batch_size / radon_half_time) print("Speedup:", astra_time / radon_time) print("Speedup half-precision:", astra_time / radon_half_time) print() title = f"Image size {args.image_size}x{args.image_size}, {args.angles} angles and batch size {args.batch_size} on a {torch.cuda.get_device_name(0)}" plot(tasks, astra_fps, radon_fps, radon_half_fps, title) if args.output: plt.savefig(args.output, dpi=300) else: plt.show()
class Operators(): def __init__(self, img_size, n_angles, sample_ratio, device, circle=False, beam='parallel', det_size=None, det_dist=(1800, 0)): self.device = device self.img_size = img_size self.sample_ratio = sample_ratio self.n_angles = n_angles self.det_size = img_size if det_size == None else det_size self.init_radon(beam, circle, det_dist) self.landweber = Landweber(self.radon) self.mask = torch.zeros((1, 1, 1, n_angles)).to(device) self.mask[:, :, :, ::sample_ratio].fill_(1) def init_radon(self, beam, circle, det_dist): if beam == 'parallel': angles = np.linspace(0, np.pi, self.n_angles, endpoint=False) self.radon = Radon(self.img_size, angles, clip_to_circle=circle) self.radon_sparse = Radon(self.img_size, angles[::self.sample_ratio], clip_to_circle=circle) elif beam == 'fan': angles = np.linspace(0, self.n_angles / 180 * np.pi, self.n_angles, False) self.radon = RadonFanbeam(self.img_size, angles, source_distance=det_dist[0], det_distance=det_dist[1], clip_to_circle=circle, det_count=self.det_size) self.radon_sparse = RadonFanbeam(self.img_size, angles[::self.sample_ratio], source_distance=det_dist[0], det_distance=det_dist[1], clip_to_circle=circle, det_count=self.det_size) else: raise Exception('projection beam type undefined!') self.n_angles_sparse = len(angles[::self.sample_ratio]) # Radon Transform. def forward_radon(self, input): # check dimension if list(input.shape[2:]) != [self.img_size, self.img_size]: raise Exception( f'radon input dimension wrong! received {input.size()}.') return self.radon.forward(input) / self.det_size # $X^\T ()$ inverse radon def forward_adjoint(self, input): # check dimension if list(input.shape[2:]) == [self.n_angles, self.det_size]: return self.radon.backprojection(input) / self.n_angles elif list(input.shape[2:]) == [self.n_angles_sparse, self.det_size]: return self.radon_sparse.backprojection( input) / self.n_angles_sparse # scale the angles else: raise Exception( f'forward_adjoint input dimension wrong! received {input.size()}.' ) # $X^\T X ()$ def forward_gramian(self, input): sinogram = self.forward_radon(input) return self.forward_adjoint(sinogram) # Corruption model: undersample sinogram by sample_ratio def undersample_model(self, input): return input[:, :, ::self.sample_ratio, :] # estimate step size eta def estimate_eta(self): eta = self.landweber.estimate_alpha(self.img_size, self.device) return torch.tensor(eta, dtype=torch.float32, device=self.device)
proj_geom = astra.create_proj_geom('fanflat', det_width, x.shape[0], angles, source_distance, det_distance) #proj_geom = astra.create_proj_geom('parallel', det_width, x.shape[0], angles) proj_id = astra.create_projector('cuda', proj_geom, vol_geom) id, astra_y = astra.create_sino(x, proj_id) _, astra_bp = astra.create_backprojection(astra_y, proj_id) plt.imshow(astra_bp) # # # # plt.figure() # # # projection = FanBeamProjection(x.shape[0], source_distance, det_distance, det_width) # # # projection = ParallelBeamProjection(x.shape[0]) radon = RadonFanbeam(x.shape[-1], angles, source_distance, det_distance, det_spacing=det_width) print(source_distance + det_distance) # radon.rays = fan_beam_rays(x.shape[0], source_distance, det_distance, det_width, clip_to_circle=True) # print(radon.rays.size()) x = torch.FloatTensor(x).cuda().view(1, x.shape[0], x.shape[1]) # repeat data to fill batch size x = torch.cat([x] * 8, dim=0) y = radon.forward(x) bp = radon.backprojection(y) print(np.linalg.norm(astra_y - y.cpu().numpy()) / np.linalg.norm(astra_y)) scale = np.mean(astra_bp) / np.mean(bp[0].cpu().numpy()) print(scale)
def main(): parser = argparse.ArgumentParser(description='Benchmark and compare with Astra Toolbox') parser.add_argument('--task', default="all") parser.add_argument('--image-size', default=256, type=int) parser.add_argument('--angles', default=-1, type=int) parser.add_argument('--batch-size', default=32, type=int) parser.add_argument('--samples', default=50, type=int) parser.add_argument('--warmup', default=10, type=int) parser.add_argument('--output', default="") parser.add_argument('--circle', action='store_true') args = parser.parse_args() if args.angles == -1: args.angles = args.image_size device = torch.device("cuda") angles = np.linspace(0, 2 * np.pi, args.angles, endpoint=False).astype(np.float32) radon = Radon(args.image_size, angles, clip_to_circle=args.circle) radon_fb = RadonFanbeam(args.image_size, angles, args.image_size, clip_to_circle=args.circle) astra = AstraWrapper(angles) if args.task == "all": tasks = ["forward", "backward", "fanbeam forward", "fanbeam backward"] else: tasks = [args.task] astra_fps = [] radon_fps = [] radon_half_fps = [] # x = torch.randn((args.batch_size, args.image_size, args.image_size), device=device) # if "forward" in tasks: # print("Benchmarking forward") # x = generate_random_images(args.batch_size, args.image_size) # astra_time = benchmark_function(lambda y: astra.forward(y), x, args.samples, args.warmup) # radon_time = benchmark_function(lambda y: radon.forward(torch.FloatTensor(x).to(device)).cpu(), x, args.samples, # args.warmup) # radon_half_time = benchmark_function(lambda y: radon.forward(torch.HalfTensor(x).to(device)).cpu(), x, # args.samples, args.warmup) # # astra_fps.append(args.batch_size / astra_time) # radon_fps.append(args.batch_size / radon_time) # radon_half_fps.append(args.batch_size / radon_half_time) # # print(astra_time, radon_time, radon_half_time) # astra.clean() # # if "backward" in tasks: # print("Benchmarking backward") # x = generate_random_images(args.batch_size, args.image_size) # pid, x = astra.forward(x) # # astra_time = benchmark_function(lambda y: astra.backproject(pid, args.image_size, args.batch_size), x, # args.samples, args.warmup) # radon_time = benchmark_function(lambda y: radon.backward(torch.FloatTensor(x).to(device)).cpu(), x, # args.samples, # args.warmup) # radon_half_time = benchmark_function(lambda y: radon.backward(torch.HalfTensor(x).to(device)).cpu(), x, # args.samples, args.warmup) # # astra_fps.append(args.batch_size / astra_time) # radon_fps.append(args.batch_size / radon_time) # radon_half_fps.append(args.batch_size / radon_half_time) # # print(astra_time, radon_time, radon_half_time) # astra.clean() # if "forward+backward" in tasks: # print("Benchmarking forward + backward") # x = generate_random_images(args.batch_size, args.image_size) # astra_time = benchmark_function(lambda y: astra_forward_backward(astra, y, args.image_size, args.batch_size), x, # args.samples, args.warmup) # radon_time = benchmark_function(lambda y: radon_forward_backward(radon, y), x, args.samples, # args.warmup) # radon_half_time = benchmark_function(lambda y: radon_forward_backward(radon, y, half=True), x, # args.samples, args.warmup) # astra_fps.append(args.batch_size / astra_time) # radon_fps.append(args.batch_size / radon_time) # radon_half_fps.append(args.batch_size / radon_half_time) # print(astra_time, radon_time, radon_half_time) # astra.clean() if "forward" in tasks: print("Benchmarking forward from device") x = generate_random_images(args.batch_size, args.image_size) dx = torch.FloatTensor(x).to(device) astra_time = benchmark_function(lambda y: astra.forward(y), x, args.samples, args.warmup) radon_time = benchmark_function(lambda y: radon.forward(y), dx, args.samples, args.warmup, sync=True) radon_half_time = benchmark_function(lambda y: radon.forward(y), dx.half(), args.samples, args.warmup, sync=True) astra_fps.append(args.batch_size / astra_time) radon_fps.append(args.batch_size / radon_time) radon_half_fps.append(args.batch_size / radon_half_time) print(astra_time, radon_time, radon_half_time) astra.clean() if "backward" in tasks: print("Benchmarking backward from device") x = generate_random_images(args.batch_size, args.image_size) dx = torch.FloatTensor(x).to(device) pid, x = astra.forward(x) astra_time = benchmark_function(lambda y: astra.backproject(pid, args.image_size, args.batch_size), x, args.samples, args.warmup) radon_time = benchmark_function(lambda y: radon.backward(y), dx, args.samples, args.warmup, sync=True) radon_half_time = benchmark_function(lambda y: radon.backward(y), dx.half(), args.samples, args.warmup, sync=True) astra_fps.append(args.batch_size / astra_time) radon_fps.append(args.batch_size / radon_time) radon_half_fps.append(args.batch_size / radon_half_time) print(astra_time, radon_time, radon_half_time) astra.clean() if "fanbeam forward" in tasks: print("Benchmarking fanbeam forward") x = generate_random_images(args.batch_size, args.image_size) dx = torch.FloatTensor(x).to(device) # # astra_time = benchmark_function(lambda y: astra.backproject(pid, args.image_size, args.batch_size), x, # args.samples, args.warmup) radon_time = benchmark_function(lambda y: radon_fb.forward(y), dx, args.samples, args.warmup, sync=True) radon_half_time = benchmark_function(lambda y: radon_fb.forward(y), dx.half(), args.samples, args.warmup, sync=True) astra_fps.append(0.0) radon_fps.append(args.batch_size / radon_time) radon_half_fps.append(args.batch_size / radon_half_time) #print(astra_time, radon_time, radon_half_time) astra.clean() if "fanbeam backward" in tasks: print("Benchmarking fanbeam backward") x = generate_random_images(args.batch_size, args.image_size) dx = torch.FloatTensor(x).to(device) # # astra_time = benchmark_function(lambda y: astra.backproject(pid, args.image_size, args.batch_size), x, # args.samples, args.warmup) radon_time = benchmark_function(lambda y: radon_fb.backprojection(y), dx, args.samples, args.warmup, sync=True) radon_half_time = benchmark_function(lambda y: radon_fb.backprojection(y), dx.half(), args.samples, args.warmup, sync=True) astra_fps.append(0.0) radon_fps.append(args.batch_size / radon_time) radon_half_fps.append(args.batch_size / radon_half_time) #print(astra_time, radon_time, radon_half_time) astra.clean() title = f"Image size {args.image_size}x{args.image_size}, {args.angles} angles and batch size {args.batch_size} on a {torch.cuda.get_device_name(0)}" plot(tasks, astra_fps, radon_fps, radon_half_fps, title) if args.output: plt.savefig(args.output, dpi=300) else: plt.show()