https://github.com/jyhmiinlin/pynufft
Tip revision: fd2bca8beca742517794bef0efb1f5014e494586 authored by Jyh-Miin Lin on 26 September 2020, 04:17:59 UTC
change the version to 2020.1.2
change the version to 2020.1.2
Tip revision: fd2bca8
test_cuda_2D.py
"""
Explicitly load the NUFFT_hsa to the 'cuda' backend.
"""
from pynufft import NUFFT_cpu, NUFFT_hsa
import numpy
def test_cuda():
import numpy
import matplotlib.pyplot
# load example image
import pkg_resources
## Define the source of data
DATA_PATH = pkg_resources.resource_filename('pynufft', 'src/data/')
# PHANTOM_FILE = pkg_resources.resource_filename('pynufft', 'data/phantom_256_256.txt')
import scipy
image = scipy.misc.ascent()[::2,::2]
image=image.astype(numpy.float)/numpy.max(image[...])
Nd = (256, 256) # image space size
Kd = (512, 512) # k-space size
Jd = (6, 6) # interpolation size
# load k-space points as M * 2 array
om = numpy.load(DATA_PATH+'om2D.npz')['arr_0']
# Show the shape of om
print('the shape of om = ', om.shape)
# initiating NUFFT_cpu object
nfft = NUFFT_cpu() # CPU NUFFT class
# Plan the nfft object
nfft.plan(om, Nd, Kd, Jd)
# initiating NUFFT_hsa object
NufftObj = NUFFT_hsa('cuda', 0, 0)
# Plan the NufftObj (similar to NUFFT_cpu)
NufftObj.plan(om, Nd, Kd, Jd)
import time
t0 = time.time()
for pp in range(0,10):
y = nfft.forward(image)
t_cpu = (time.time() - t0)/10.0
## Moving image to gpu
## gx is an gpu array, dtype = complex64
gx = NufftObj.to_device(image)
t0= time.time()
for pp in range(0,100):
gy = NufftObj.forward(gx)
t_cu = (time.time() - t0)/100
print('t_cpu = ', t_cpu)
print('t_cuda =, ', t_cu)
print('gy close? = ', numpy.allclose(y, gy.get(), atol=numpy.linalg.norm(y)*1e-3))
print("acceleration=", t_cpu/t_cu)
maxiter =100
import time
t0= time.time()
x_cpu_cg = nfft.solve(y, 'cg',maxiter=maxiter)
# x2 = nfft.solve(y2, 'L1TVLAD',maxiter=maxiter, rho = 2)
t1 = time.time()-t0
# gy=NufftObj.thr.copy_array(NufftObj.thr.to_device(y2))
t0= time.time()
x_cuda_cg = NufftObj.solve(gy,'cg', maxiter=maxiter)
# x = NufftObj.solve(gy,'L1TVLAD', maxiter=maxiter, rho=2)
t2 = time.time() - t0
print(t1, t2)
print('acceleration of cg=', t1/t2 )
t0= time.time()
x_cpu_TV = nfft.solve(y, 'L1TVOLS',maxiter=maxiter, rho = 2)
t1 = time.time()-t0
t0= time.time()
x_cuda_TV = NufftObj.solve(gy,'L1TVOLS', maxiter=maxiter, rho=2)
t2 = time.time() - t0
print(t1, t2)
print('acceleration of TV=', t1/t2 )
matplotlib.pyplot.subplot(2, 2, 1)
matplotlib.pyplot.imshow( x_cpu_cg.real, cmap= matplotlib.cm.gray)
matplotlib.pyplot.title('CG_cpu')
matplotlib.pyplot.subplot(2, 2, 2)
matplotlib.pyplot.imshow(x_cuda_cg.get().real, cmap= matplotlib.cm.gray)
matplotlib.pyplot.title('CG_cuda')
matplotlib.pyplot.subplot(2, 2, 3)
matplotlib.pyplot.imshow( x_cpu_TV.real, cmap= matplotlib.cm.gray)
matplotlib.pyplot.title('TV_cpu')
matplotlib.pyplot.subplot(2, 2, 4)
matplotlib.pyplot.imshow(x_cuda_TV.get().real, cmap= matplotlib.cm.gray)
matplotlib.pyplot.title('TV_cuda')
matplotlib.pyplot.show()
NufftObj.release()
del NufftObj
if __name__ == '__main__':
test_cuda()