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 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()
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)
f = torch.zeros_like(bp)
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()
# 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)
plt.plot(values)
s = time.time()
for _ in range(1):
with torch.no_grad():
rec, values = cgne(
radon,