def test_half(device, batch_size, image_size, angles, spacing, 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) # our implementation radon = Radon(image_size, angles, det_spacing=spacing, det_count=det_count, clip_to_circle=clip_to_circle) x = torch.FloatTensor(x).to(device) sinogram = radon.forward(x) single_precision = radon.backprojection(sinogram) h_sino = radon.forward(x.half()) half_precision = radon.backprojection(h_sino) 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(): n_angles = 100 image_size = 512 circle_radius = 100 source_dist = 1.5 * image_size batch_size = 1 n_scales = 5 angles = (np.linspace(0., 100., n_angles, endpoint=False) - 50.0) / 180.0 * np.pi x = np.zeros((image_size, image_size), dtype=np.float32) x[circle_mask(image_size, circle_radius)] = 1.0 radon = Radon(image_size, angles) # RadonFanbeam(image_size, angles, source_dist) shearlet = ShearletTransform(image_size, image_size, [0.5] * n_scales) torch_x = torch.from_numpy(x).cuda() torch_x = torch_x.view(1, image_size, image_size).repeat(batch_size, 1, 1) sinogram = radon.forward(torch_x) bp = radon.backward(sinogram) sc = shearlet.forward(bp) p_0 = 0.02 p_1 = 0.1 w = 3**shearlet.scales / 400 w = w.view(1, -1, 1, 1).cuda() u_2 = torch.zeros_like(bp) z_2 = torch.zeros_like(bp) u_1 = torch.zeros_like(sc) z_1 = torch.zeros_like(sc) f = torch.zeros_like(bp) relative_error = [] start_time = time.time() for i in range(100): cg_y = p_0 * bp + p_1 * shearlet.backward(z_1 - u_1) + (z_2 - u_2) f = cg(lambda x: p_0 * radon.backward(radon.forward(x)) + (1 + p_1) * x, f.clone(), cg_y, max_iter=50) sh_f = shearlet.forward(f) z_1 = shrink(sh_f + u_1, p_0 / p_1 * w) z_2 = (f + u_2).clamp_min(0) u_1 = u_1 + sh_f - z_1 u_2 = u_2 + f - z_2 relative_error.append( (torch.norm(torch_x[0] - f[0]) / torch.norm(torch_x[0])).item()) runtime = time.time() - start_time print("Running time:", runtime) print("Running time per image:", runtime / batch_size) print("Relative error: ", 100 * relative_error[-1])
def bench_parallel_backward(phantom, det_count, num_angles, warmup, repeats): radon = Radon(phantom.size(1), np.linspace(0, np.pi, num_angles, endpoint=False), det_count) sino = radon.forward(phantom) f = lambda x: radon.forward(x) return benchmark(f, sino, warmup, repeats)
def test_error(device, batch_size, image_size, angles, spacing, clip_to_circle): # generate random images x = generate_random_images(batch_size, image_size, masked=clip_to_circle) # astra astra = AstraWrapper(angles) astra_fp_id, astra_fp = astra.forward(x, spacing) astra_bp = astra.backproject(astra_fp_id, image_size, batch_size) if clip_to_circle: astra_bp *= circle_mask(image_size) # our implementation radon = Radon(image_size, angles, det_spacing=spacing, clip_to_circle=clip_to_circle) x = torch.FloatTensor(x).to(device) our_fp = radon.forward(x) our_bp = radon.backprojection(our_fp) forward_error = relative_error(astra_fp, our_fp.cpu().numpy()) back_error = relative_error(astra_bp, our_bp.cpu().numpy()) # if forward_error > 10: # plt.imshow(astra_fp[0]) # plt.figure() # plt.imshow(our_fp[0].cpu().numpy()) # plt.show() print( f"batch: {batch_size}, size: {image_size}, angles: {len(angles)}, spacing: {spacing}, 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 bench_fanbeam_backward(task, dtype, device, *bench_args): num_angles = task["num_angles"] det_count = task["det_count"] source_dist = task["source_distance"] det_dist = task["detector_distance"] det_spacing = task["det_spacing"] x = torch.randn(task["batch_size"], task["size"], task["size"], dtype=dtype, device=device) angles = np.linspace(0, np.pi, num_angles, endpoint=False) projection = Projection.fanbeam(source_dist, det_dist, det_count, det_spacing) radon = Radon(angles, task["size"], projection) # radon = RadonFanbeam(phantom.size(1), angles, source_dist, det_dist, det_count) sino = radon.forward(x) def f(x): return radon.backward(x) return benchmark(f, x, *bench_args)
class Operators(): def __init__(self, image_size, n_angles, sample_ratio, device, circle=False): self.device = device self.image_size = image_size self.sample_ratio = sample_ratio self.n_angles = n_angles angles = np.linspace(0, np.pi, self.n_angles, endpoint=False) self.radon = Radon(self.image_size, angles, clip_to_circle=circle) self.radon_sparse = Radon(self.image_size, angles[::sample_ratio], clip_to_circle=circle) self.n_angles_sparse = len(angles[::sample_ratio]) self.landweber = Landweber(self.radon) self.mask = torch.zeros((1,1,1,180)).to(device) self.mask[:,:,:,::sample_ratio].fill_(1) # $X^\T ()$ inverse radon def forward_adjoint(self, input): # check dimension if input.size()[3] == self.n_angles: return self.radon.backprojection(input.permute(0,1,3,2)) elif input.size()[3] == self.n_angles_sparse: return self.radon_sparse.backprojection(input.permute(0,1,3,2))/self.n_angles_sparse*self.n_angles # scale the angles else: raise Exception(f'forward_adjoint input dimension wrong! received {input.size()}.') # $X^\T X ()$ def forward_gramian(self, input): # check dimension if input.size()[2] != self.image_size: raise Exception(f'forward_gramian input dimension wrong! received {input.size()}.') sinogram = self.radon.forward(input) return self.radon.backprojection(sinogram) # Corruption model: undersample sinogram by 8 def undersample_model(self, input): return input[:,:,:,::self.sample_ratio] # Filtered Backprojection. Input siogram range = (0,1) def FBP(self, input): # check dimension if input.size()[2] != self.image_size or input.size()[3] != self.n_angles: raise Exception(f'FBP input dimension wrong! received {input.size()}.') filtered_sinogram = self.radon.filter_sinogram(input.permute(0,1,3,2)) return self.radon.backprojection(filtered_sinogram) # estimate step size eta def estimate_eta(self): eta = self.landweber.estimate_alpha(self.image_size, self.device) return torch.tensor(eta, dtype=torch.float32, device=self.device)
def test_shapes(self): """ Check using channels is ok """ device = torch.device('cuda') angles = torch.FloatTensor( np.linspace(0, 2 * np.pi, 10).astype(np.float32)).to(device) radon = Radon(64, angles) # test with 2 batch dimensions x = torch.FloatTensor(2, 3, 64, 64).to(device) y = radon.forward(x) self.assertEqual(y.size(), (2, 3, 10, 64)) z = radon.backprojection(y) self.assertEqual(z.size(), (2, 3, 64, 64)) # no batch dimensions x = torch.FloatTensor(64, 64).to(device) y = radon.forward(x) self.assertEqual(y.size(), (10, 64)) z = radon.backprojection(y) self.assertEqual(z.size(), (64, 64))
def test_differentiation(self): device = torch.device('cuda') x = torch.FloatTensor(1, 64, 64).to(device) x.requires_grad = True angles = torch.FloatTensor( np.linspace(0, 2 * np.pi, 10).astype(np.float32)).to(device) radon = Radon(64, angles) # check that backward is implemented for fp and bp y = radon.forward(x) z = torch.mean(radon.backprojection(y)) z.backward() self.assertIsNotNone(x.grad)
def test_noise(): device = torch.device('cuda') x = torch.FloatTensor(3, 5, 64, 64).to(device) lookup_table = torch.FloatTensor(128, 64).to(device) x.requires_grad = True angles = torch.FloatTensor(np.linspace(0, 2 * np.pi, 10).astype(np.float32)) radon = Radon(64, angles) sinogram = radon.forward(x) assert_equal(sinogram.size(), (3, 5, 10, 64)) readings = radon.emulate_readings(sinogram, 5, 10.0) assert_equal(readings.size(), (3, 5, 10, 64)) assert_equal(readings.dtype, torch.int32) y = radon.readings_lookup(readings, lookup_table) assert_equal(y.size(), (3, 5, 10, 64)) assert_equal(y.dtype, torch.float32)
def bench_parallel_backward(task, dtype, device, *bench_args): num_angles = task["num_angles"] det_count = task["det_count"] x = torch.randn(task["batch_size"], task["size"], task["size"], dtype=dtype, device=device) angles = np.linspace(0, np.pi, num_angles, endpoint=False) projection = Projection.parallel_beam(det_count) radon = Radon(angles, task["size"], projection) # radon = Radon(phantom.size(1), np.linspace(0, np.pi, num_angles, endpoint=False), det_count) sino = radon.forward(x) def f(x): return radon.backward(x) return benchmark(f, sino, *bench_args)
channels = 4 device = torch.device('cuda') criterion = nn.L1Loss() # Instantiate a model for the sinogram and one for the image sino_model = nn.Conv2d(1, channels, 5, padding=2).to(device) image_model = nn.Conv2d(channels, 1, 3, padding=1).to(device) # create empty images x = torch.FloatTensor(batch_size, 1, image_size, image_size).to(device) # instantiate Radon transform angles = np.linspace(0, np.pi, n_angles) radon = Radon(image_size, angles) # forward projection sinogram = radon.forward(x) # apply sino_model to sinograms filtered_sinogram = sino_model(sinogram) # backprojection backprojected = radon.backprojection(filtered_sinogram) # apply image_model to backprojected images y = image_model(backprojected) # backward works as usual loss = criterion(y, x) loss.backward()
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()
# rdx *= (alpha_e - alpha_s) # rdy *= (alpha_e - alpha_s) # # print(rsx, rsy, rsx**2 + rsy**2 - v**2) # print(rdx, rdy, (rsx+rdx)**2 + (rsy+rdy)**2 - v**2) device = torch.device('cuda') angles = np.linspace(0, 2 * np.pi, 180).astype(np.float32) batch_size = 4 image_size = 256 astraw = AstraWrapper(angles) x = generate_random_images(batch_size, image_size, masked=True) astra_fp_id, astra_fp = astraw.forward(x) # our implementation radon = Radon(image_size, angles, clip_to_circle=True) x = torch.FloatTensor(x).to(device) our_fp = radon.forward(x) plt.imshow(astra_fp[0]) plt.figure() plt.imshow(our_fp[0].cpu().numpy()) plt.show() print(relative_error(astra_fp, our_fp.cpu().numpy()))
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()
batch_size = 1 n_angles = 512 // 4 image_size = 512 img = np.load("phantom.npy") device = torch.device('cuda') # instantiate Radon transform angles = np.linspace(0, np.pi / 4, n_angles, endpoint=False) radon = Radon(image_size, angles) shearlet = Shearlet(512, 512, [0.5] * 5, cache=None) # ".cache") with torch.no_grad(): x = torch.FloatTensor(img).reshape(1, image_size, image_size).to(device) sinogram = radon.forward(x) bp = radon.backward(sinogram, extend=False) # f, values = CG(radon, 1.0 / 512**2, 0.0001, bp.clone(), bp) # # print(torch.norm(x - f)/torch.norm(x)) sc = shearlet.forward(bp) p_0 = 0.02 p_1 = 0.1 w = 3**shearlet.scales / 400 w = w.view(1, -1, 1, 1).to(device) u_2 = torch.zeros_like(bp) z_2 = torch.zeros_like(bp) u_1 = torch.zeros_like(sc) z_1 = torch.zeros_like(sc)
batch_size = 1 n_angles = 512 image_size = 512 img = np.load("phantom.npy") device = torch.device('cuda') # instantiate Radon transform angles = np.linspace(0, np.pi, n_angles, endpoint=False) radon = Radon(image_size, angles) x = torch.FloatTensor(img).to(device).view(1, 512, 512) x = torch.cat([x] * 4, dim=0).view(2, 2, 512, 512) print(x.size()) y = radon.forward(x) # CG(radon, 1.0, 0.0, torch.zeros_like(x), radon.backward(y)) # rec = cgne(radon, torch.zeros_like(x), y, tol=1e-2) s = time.time() for _ in range(1): with torch.no_grad(): rec, values = cg(lambda z: radon.backward(radon.forward(z)), torch.zeros_like(x), radon.backward(y), callback=lambda x, r: torch.norm( radon.forward(x[0]) - y[0]).item(), max_iter=500) print(torch.norm(x - rec).item() / torch.norm(x).item()) print("CG", time.time() - s)