Note
Go to the end to download the full example code.
pytorch parallelproj projection layer¶
In this example, we show how to define a custom pytorch layer that can be used
to define a feed forward neural network that includes a parallelproj forward and back
backward projections (or any LinearOperator) that can be used with pytorch’s
autograd engine.
from __future__ import annotations
import array_api_compat.torch as torch
import matplotlib.pyplot as plt
import parallelproj.operators
import parallelproj.projectors
import parallelproj.pet_scanners
import parallelproj.pet_lors
import parallelproj_core as ppc
from array_api_compat import device
# device variable (cpu or cuda) that determines whether calculations
# are performed on the cpu or cuda gpu
if torch.cuda.is_available() and ppc.cuda_enabled == 1:
dev = "cuda"
else:
dev = "cpu"
Setup the forward projection layer¶
We subclass torch.autograd.Function to define a custom pytorch layer
that is compatible with pytorch’s autograd engine.
see also: https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html
class LinearSingleChannelOperator(torch.autograd.Function):
"""
Function representing a linear operator acting on a mini batch of single channel images
"""
@staticmethod
def forward(
ctx, x: torch.Tensor, operator: parallelproj.operators.LinearOperator
) -> torch.Tensor:
"""forward pass of the linear operator
Parameters
----------
ctx : context object
that can be used to store information for the backward pass
x : torch.Tensor
mini batch of 3D images with dimension (batch_size, 1, num_voxels_x, num_voxels_y, num_voxels_z)
operator : parallelproj.operators.LinearOperator
linear operator that can act on a single 3D image
Returns
-------
torch.Tensor
mini batch of 3D images with dimension (batch_size, *operator.out_shape)
"""
# https://pytorch.org/docs/stable/notes/extending.html#how-to-use
ctx.set_materialize_grads(False)
ctx.operator = operator
batch_size = x.shape[0]
y = torch.zeros(
(batch_size,) + operator.out_shape, dtype=x.dtype, device=device(x)
)
# loop over all samples in the batch and apply linear operator
# to the first channel
for i in range(batch_size):
y[i, ...] = operator(x[i, 0, ...].detach())
return y
@staticmethod
def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None]:
"""backward pass of the forward pass
Parameters
----------
ctx : context object
that can be used to obtain information from the forward pass
grad_output : torch.Tensor
mini batch of dimension (batch_size, *operator.out_shape)
Returns
-------
torch.Tensor, None
mini batch of 3D images with dimension (batch_size, 1, *operator.in_shape)
"""
# For details on how to implement the backward pass, see
# https://pytorch.org/docs/stable/notes/extending.html#how-to-use
# since forward takes two input arguments (x, operator)
# we have to return two arguments (the latter is None)
if grad_output is None:
return None, None
else:
operator = ctx.operator
batch_size = grad_output.shape[0]
x = torch.zeros(
(batch_size, 1) + operator.in_shape,
dtype=grad_output.dtype,
device=device(grad_output),
)
# loop over all samples in the batch and apply linear operator
# to the first channel
for i in range(batch_size):
x[i, 0, ...] = operator.adjoint(grad_output[i, ...].detach())
return x, None
Setup the back projection layer¶
We subclass torch.autograd.Function to define a custom pytorch layer
that is compatible with pytorch’s autograd engine.
see also: https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html
class AdjointLinearSingleChannelOperator(torch.autograd.Function):
"""
Function representing the adjoint of a linear operator acting on a mini batch of single channel images
"""
@staticmethod
def forward(
ctx, x: torch.Tensor, operator: parallelproj.operators.LinearOperator
) -> torch.Tensor:
"""forward pass of the adjoint of the linear operator
Parameters
----------
ctx : context object
that can be used to store information for the backward pass
x : torch.Tensor
mini batch of 3D images with dimension (batch_size, *operator.out_shape)
operator : parallelproj.operators.LinearOperator
linear operator that can act on a single 3D image
Returns
-------
torch.Tensor
mini batch of 3D images with dimension (batch_size, 1, *operator.in_shape)
"""
ctx.set_materialize_grads(False)
ctx.operator = operator
batch_size = x.shape[0]
y = torch.zeros(
(batch_size, 1) + operator.in_shape, dtype=x.dtype, device=device(x)
)
# loop over all samples in the batch and apply linear operator
# to the first channel
for i in range(batch_size):
y[i, 0, ...] = operator.adjoint(x[i, ...].detach())
return y
@staticmethod
def backward(ctx, grad_output):
"""backward pass of the forward pass
Parameters
----------
ctx : context object
that can be used to obtain information from the forward pass
grad_output : torch.Tensor
mini batch of dimension (batch_size, 1, *operator.in_shape)
Returns
-------
torch.Tensor, None
mini batch of 3D images with dimension (batch_size, *operator.out_shape)
"""
# For details on how to implement the backward pass, see
# https://pytorch.org/docs/stable/notes/extending.html#how-to-use
# since forward takes two input arguments (x, operator)
# we have to return two arguments (the latter is None)
if grad_output is None:
return None, None
else:
operator = ctx.operator
batch_size = grad_output.shape[0]
x = torch.zeros(
(batch_size,) + operator.out_shape,
dtype=grad_output.dtype,
device=device(grad_output),
)
# loop over all samples in the batch and apply linear operator
# to the first channel
for i in range(batch_size):
x[i, ...] = operator(grad_output[i, 0, ...].detach())
return x, None
Setup a minimal non-TOF PET projector¶
We setup a minimal non-TOF PET projector of small scanner with three rings.
num_rings = 3
scanner = parallelproj.pet_scanners.RegularPolygonPETScannerGeometry(
torch,
dev,
radius=35.0,
num_sides=12,
num_lor_endpoints_per_side=6,
lor_spacing=3.0,
ring_positions=torch.linspace(-4, 4, num_rings, device=dev),
symmetry_axis=2,
)
# setup the LOR descriptor that defines the sinogram
lor_desc = parallelproj.pet_lors.RegularPolygonPETLORDescriptor(
scanner,
parallelproj.pet_lors.Michelogram(scanner.num_rings, max_ring_difference=1, span=1),
radial_trim=10,
sinogram_order=parallelproj.pet_lors.SinogramSpatialAxisOrder.RVP,
)
proj = parallelproj.projectors.RegularPolygonPETProjector(
lor_desc, img_shape=(20, 20, 5), voxel_size=(2.0, 2.0, 2.0)
)
Define a mini batch of input and output tensors¶
batch_size = 2
x = torch.rand(
(batch_size, 1) + proj.in_shape,
device=dev,
dtype=torch.float32,
requires_grad=True,
)
y = torch.rand(
(batch_size,) + proj.out_shape,
device=dev,
dtype=torch.float32,
requires_grad=True,
)
Define the forward and backward projection layers¶
fwd_op_layer = LinearSingleChannelOperator.apply
adjoint_op_layer = AdjointLinearSingleChannelOperator.apply
f1 = fwd_op_layer(x, proj)
print("forward projection (Ax) .:", f1.shape, type(f1), device(f1))
b1 = adjoint_op_layer(y, proj)
print("back projection (A^T y) .:", b1.shape, type(b1), device(b1))
fb1 = adjoint_op_layer(fwd_op_layer(x, proj), proj)
print("back + forward projection (A^TAx) .:", fb1.shape, type(fb1), device(fb1))
forward projection (Ax) .: torch.Size([2, 51, 36, 7]) <class 'torch.Tensor'> cpu
back projection (A^T y) .: torch.Size([2, 1, 20, 20, 5]) <class 'torch.Tensor'> cpu
back + forward projection (A^TAx) .: torch.Size([2, 1, 20, 20, 5]) <class 'torch.Tensor'> cpu
Define a dummy loss function and trigger the backpropagation¶
# define a dummy loss function
dummy_loss = (fb1**2).sum()
# trigger the backpropagation
dummy_loss.backward()
print(f" backpropagted gradient {x.grad}")
backpropagted gradient tensor([[[[[16932896.0000, 6352481.0000, 27397480.0000, 6380920.5000,
17273928.0000],
[19099642.0000, 8492678.0000, 30550182.0000, 8595404.0000,
19520214.0000],
[20724762.0000, 10817195.0000, 32533588.0000, 10969943.0000,
21175812.0000],
...,
[21330666.0000, 11219816.0000, 32817322.0000, 10795242.0000,
21046446.0000],
[19716474.0000, 8825927.0000, 30585234.0000, 8521366.0000,
19409140.0000],
[17443932.0000, 6524830.5000, 27436158.0000, 6335062.5000,
17214780.0000]],
[[19079140.0000, 8518590.0000, 30548918.0000, 8635200.0000,
19515670.0000],
[21970598.0000, 11735578.0000, 34764804.0000, 11939253.0000,
22518344.0000],
[23001820.0000, 14203363.0000, 35754420.0000, 14447078.0000,
23564554.0000],
...,
[23655534.0000, 14826847.0000, 35964320.0000, 14242244.0000,
23373646.0000],
[22714170.0000, 12212977.0000, 34842596.0000, 11741427.0000,
22380786.0000],
[19671780.0000, 8783159.0000, 30601994.0000, 8515551.0000,
19473182.0000]],
[[20694072.0000, 10785287.0000, 32655780.0000, 10952261.0000,
21163144.0000],
[22953230.0000, 14239013.0000, 35821944.0000, 14495623.0000,
23553086.0000],
[23803294.0000, 17099806.0000, 36531948.0000, 17369572.0000,
24395942.0000],
...,
[24420630.0000, 17774280.0000, 36707048.0000, 17094096.0000,
24189944.0000],
[23678734.0000, 14788140.0000, 35898872.0000, 14172911.0000,
23389088.0000],
[21287292.0000, 11221167.0000, 32686326.0000, 10809938.0000,
21046232.0000]],
...,
[[21217240.0000, 11126856.0000, 32737468.0000, 10992263.0000,
21453314.0000],
[23564412.0000, 14672806.0000, 35934688.0000, 14433272.0000,
23863818.0000],
[24334654.0000, 17679062.0000, 36652464.0000, 17281730.0000,
24642454.0000],
...,
[24431972.0000, 17777002.0000, 36866788.0000, 17235558.0000,
24435248.0000],
[23619254.0000, 14858222.0000, 36209088.0000, 14413261.0000,
23643350.0000],
[21265478.0000, 11235368.0000, 33006884.0000, 10911495.0000,
21265736.0000]],
[[19561894.0000, 8698728.0000, 30661058.0000, 8625960.0000,
19766036.0000],
[22500922.0000, 12075355.0000, 34907808.0000, 11956471.0000,
22819774.0000],
[23496694.0000, 14703118.0000, 35941144.0000, 14452310.0000,
23797616.0000],
...,
[23667030.0000, 14831042.0000, 36237068.0000, 14385982.0000,
23675464.0000],
[22625018.0000, 12250281.0000, 35251120.0000, 11904508.0000,
22645544.0000],
[19619982.0000, 8857850.0000, 30965348.0000, 8628376.0000,
19625624.0000]],
[[17276590.0000, 6472267.5000, 27512382.0000, 6406755.0000,
17471192.0000],
[19455524.0000, 8706516.0000, 30630314.0000, 8664995.0000,
19763740.0000],
[21130730.0000, 11084881.0000, 32789740.0000, 10970615.0000,
21402508.0000],
...,
[21285094.0000, 11266842.0000, 32980902.0000, 10931723.0000,
21243362.0000],
[19655728.0000, 8833125.0000, 30987346.0000, 8609263.0000,
19656150.0000],
[17381372.0000, 6559932.5000, 27789476.0000, 6404095.5000,
17386040.0000]]]],
[[[[17453540.0000, 6426334.5000, 27569554.0000, 6539884.5000,
17709240.0000],
[19719024.0000, 8625850.0000, 30683532.0000, 8757020.0000,
19995464.0000],
[21380000.0000, 10990309.0000, 32706550.0000, 11160870.0000,
21679512.0000],
...,
[21413716.0000, 11182053.0000, 33643548.0000, 11302536.0000,
21275740.0000],
[19750886.0000, 8852370.0000, 31631352.0000, 8929781.0000,
19630158.0000],
[17503398.0000, 6558547.0000, 28375730.0000, 6621824.5000,
17398850.0000]],
[[19715248.0000, 8657915.0000, 30667510.0000, 8787188.0000,
19941534.0000],
[22736022.0000, 11934922.0000, 34872032.0000, 12113318.0000,
23030538.0000],
[23771188.0000, 14427895.0000, 35903972.0000, 14645279.0000,
24076510.0000],
...,
[23826300.0000, 14756789.0000, 36982544.0000, 14937423.0000,
23609876.0000],
[22828278.0000, 12230759.0000, 36047136.0000, 12324991.0000,
22587968.0000],
[19800720.0000, 8828948.0000, 31652212.0000, 8900760.0000,
19633938.0000]],
[[21362968.0000, 10956920.0000, 32753396.0000, 11112851.0000,
21604830.0000],
[23721838.0000, 14445152.0000, 35876188.0000, 14647468.0000,
23985268.0000],
[24579326.0000, 17297700.0000, 36617972.0000, 17538658.0000,
24818194.0000],
...,
[24713816.0000, 17688978.0000, 37772620.0000, 17898656.0000,
24368340.0000],
[23925148.0000, 14763799.0000, 37117420.0000, 14923081.0000,
23596978.0000],
[21512820.0000, 11253455.0000, 33730940.0000, 11332607.0000,
21236112.0000]],
...,
[[21176776.0000, 11195371.0000, 33456984.0000, 11413606.0000,
21575896.0000],
[23485108.0000, 14727279.0000, 36761996.0000, 15002947.0000,
23975026.0000],
[24239052.0000, 17684870.0000, 37508472.0000, 17975980.0000,
24795268.0000],
...,
[23795510.0000, 17448168.0000, 38071148.0000, 17875876.0000,
24311138.0000],
[22976302.0000, 14534755.0000, 37149772.0000, 14891490.0000,
23533728.0000],
[20746876.0000, 11015792.0000, 33783488.0000, 11290141.0000,
21241408.0000]],
[[19533742.0000, 8773919.0000, 31357566.0000, 8949866.0000,
19970990.0000],
[22486218.0000, 12180318.0000, 35709148.0000, 12419766.0000,
23000100.0000],
[23449020.0000, 14756201.0000, 36786152.0000, 15067909.0000,
24029238.0000],
...,
[23017316.0000, 14525668.0000, 37349160.0000, 14898541.0000,
23596344.0000],
[22007262.0000, 12004958.0000, 36136904.0000, 12329064.0000,
22636866.0000],
[19118072.0000, 8704756.0000, 31642100.0000, 8917210.0000,
19635786.0000]],
[[17291692.0000, 6506094.5000, 28078752.0000, 6641336.0000,
17696962.0000],
[19512912.0000, 8802124.0000, 31351354.0000, 8985838.0000,
19973954.0000],
[21163630.0000, 11187562.0000, 33543126.0000, 11411037.0000,
21653642.0000],
...,
[20770672.0000, 11074748.0000, 33956064.0000, 11358930.0000,
21286236.0000],
[19171008.0000, 8694352.0000, 31740994.0000, 8910645.0000,
19671166.0000],
[17031134.0000, 6478053.5000, 28424450.0000, 6625546.5000,
17450056.0000]]]]])
Check whether the gradients are calculated correctly¶
We use pytorch’s gradcheck function to check whether the implementation of the backward pass, needed to calculate the gradients, is correct.
This test can be slow which is why we only execute it on the gpu. Note that parallelproj’s projectors use single precision precision which is why we have to use a larger atol and rtol.
if dev == "cpu":
print("skipping (slow) gradient checks on cpu")
else:
print("Running forward projection layer gradient test")
grad_test_fwd = torch.autograd.gradcheck(
fwd_op_layer, (x, proj), eps=1e-1, atol=1e-3, rtol=1e-3
)
print("Running adjoint projection layer gradient test")
grad_test_fwd = torch.autograd.gradcheck(
adjoint_op_layer, (y, proj), eps=1e-1, atol=1e-3, rtol=1e-3
)
skipping (slow) gradient checks on cpu
Visualize the scanner geometry and image FOV¶
fig = plt.figure(figsize=(10, 10))
ax = fig.add_subplot(111, projection="3d")
ax.view_init(elev=-30, azim=160, roll=180, vertical_axis="y")
proj.show_geometry(ax)
fig.tight_layout()
fig.show()
Total running time of the script: (0 minutes 0.164 seconds)