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 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 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)
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)
print(
np.linalg.norm(astra_bp - bp[0].cpu().numpy()) / np.linalg.norm(astra_bp))
plt.figure()
plt.imshow(bp[0].cpu().numpy())
plt.show()
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()