#optimizer=torch.optim.AdamW([{'params':G.parameters(),'lr':1e-3},{'params':GV.parameters(),'lr':1e-3}]) scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.7, patience=6, verbose=True, min_lr=5e-5) trnFiles = os.listdir('/Shared/lss_jcb/abdul/prashant_cardiac_data/Data/d2/') sz = len(trnFiles) #data2=np.zeros((sz,1,n_select,N,N)).astype(np.complex64) rndm = random.sample(range(sz), sz) #%% nuf_ob = KbNufft(im_size=(nx, nx), norm='ortho').to(dtype) nuf_ob = nuf_ob.to(gpu) adjnuf_ob = AdjKbNufft(im_size=(nx, nx), norm='ortho').to(dtype) adjnuf_ob = adjnuf_ob.to(gpu) smapT = torch.ones((1, 1, 2, nx, nx)).cuda() nuf_ob = MriSenseNufft(im_size=(nx, nx), smap=smapT, norm='ortho').to(dtype) nuf_ob = nuf_ob.to(gpu) adjnuf_ob = AdjMriSenseNufft(im_size=(nx, nx), smap=smapT, norm='ortho').to(dtype) adjnuf_ob = adjnuf_ob.to(gpu) d2 = sio.loadmat('/Shared/lss_jcb/abdul/prashant_cardiac_data/Data/ktraj.mat')
#%% take them to gpu atbT = atbT.to(gpu) ktrajT1 = ktrajT1.to(gpu) smapT = smapT.to(gpu) ktrajT = ktrajT.to(gpu) #dcfT=dcfT.to(gpu) LT = LT.to(gpu) VT = VT.to(gpu) sbasis = sbasis.to(gpu) #%% generate atb sigma = 0.0 lam = 1e3 cgIter = 1 cgTol = 1e-15 nuf_ob = KbNufft(im_size=im_size).to(dtype) nuf_ob = nuf_ob.to(gpu) adjnuf_ob = AdjKbNufft(im_size=im_size).to(dtype) adjnuf_ob = adjnuf_ob.to(gpu) #real_mat, imag_mat = precomp_sparse_mats(ktrajT, adjnuf_ob) #interp_mats = {'real_interp_mats': real_mat, 'imag_interp_mats': imag_mat} nufft_ob1 = MriSenseNufft(im_size=im_size, smap=smapT2).to(dtype) nufft_ob1 = nufft_ob1.to(gpu) nufft_ob = MriSenseNufft(im_size=im_size, smap=smapT).to(dtype) nufft_ob = nufft_ob.to(gpu) adjnufft_ob = AdjMriSenseNufft(im_size=im_size, smap=smapT).to(dtype) adjnufft_ob = adjnufft_ob.to(gpu) #nufft_ob = MriSenseNufft(smap=smapT,im_size=im_size).to(dtype) #adjnufft_ob = AdjMriSenseNufft(smap=smapT,im_size=im_size ).to(dtype)
#%% x = shepp_logan_phantom().astype(np.complex) im_size = x.shape x = np.stack((np.real(x), np.imag(x))) # convert to tensor, unsqueeze batch and coil dimension # output size: (1, 1, 2, ny, nx) x = torch.tensor(x).unsqueeze(0).unsqueeze(0) klength = 64 ktraj = np.stack((np.zeros(64), np.linspace(-np.pi, np.pi, klength))) # convert to tensor, unsqueeze batch dimension # output size: (1, 2, klength) ktraj = torch.tensor(ktraj).unsqueeze(0) nufft_ob = KbNufft(im_size=im_size) # outputs a (1, 1, 2, klength) vector of k-space data kdata = nufft_ob(x, ktraj) #%% generate atb sigma = 0.0 lam = 1e-5 cgIter = 10 cgTol = 1e-15 A = lambda x: pt_A(x, csmT, maskT) At = lambda x: pt_At(x, csmT, maskT) bT = At(orgT) noiseT = torch.randn(bT.shape) * sigma noiseT = noiseT.to(gpu) bT = bT + noiseT
smap = np.stack((np.real(csmTrn), np.imag(csmTrn)), axis=1) smap=np.tile(smap,(NF,1,1,1,1)) smapT = torch.tensor(smap).to(dtype) # convert to tensor, unsqueeze batch dimension # output size: (1, 2, klength) ktraj = np.stack((np.real(ktraj), np.imag(ktraj)), axis=1) ktrajT = torch.tensor(ktraj).to(dtype) #%% take them to gpu kdataT=kdataT.cuda() smapT=smapT.cuda() ktrajT=ktrajT.cuda() dcfT=dcfT.cuda() #%% generate atb (Initial reconstruction) nuf_ob = KbNufft(im_size=im_size).to(dtype) nuf_ob=nuf_ob.cuda() adjnuf_ob = AdjKbNufft(im_size=im_size).to(dtype) adjnuf_ob=adjnuf_ob.cuda() real_mat, imag_mat = precomp_sparse_mats(ktrajT, adjnuf_ob) interp_mats = {'real_interp_mats': real_mat, 'imag_interp_mats': imag_mat} nufft_ob = MriSenseNufft(im_size=im_size, smap=smapT).to(dtype) nufft_ob=nufft_ob.cuda() adjnufft_ob = AdjMriSenseNufft(im_size=im_size, smap=smapT).to(dtype) adjnufft_ob=adjnufft_ob.cuda() At=lambda x: adjnufft_ob(x*dcfT,ktrajT,interp_mats) A=lambda x: nufft_ob(x,ktrajT,interp_mats) #inverse NUFFT transform -->Initial reconstruction
#GV.load_state_dict(torch.load('wtsVB-400.pt')) #GV.load_state_dict(torch.load('./PTmodels/27Oct_112451am_500ep_27oct/wts-500.pt')) #optimizer=torch.optim.SGD([{'params':G.parameters(),'lr':5e-3,'momentum':0.9}]) #optimizer=torch.optim.AdamW([{'params':G.parameters(),'lr':1e-4}]) optimizer=torch.optim.AdamW([{'params':G.parameters(),'lr':1e-4},{'params':GV.parameters(),'lr':1e-4}]) scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.7, patience=6, verbose=True, min_lr=1e-5) trnFiles=os.listdir('/Shared/lss_jcb/abdul/prashant_cardiac_data/Data/d2/') sz=len(trnFiles) #data2=np.zeros((sz,1,n_select,N,N)).astype(np.complex64) rndm=random.sample(range(sz),sz) #%% nuf_ob = KbNufft(im_size=(nx,nx)).to(dtype) nuf_ob=nuf_ob.to(gpu) adjnuf_ob = AdjKbNufft(im_size=(nx,nx)).to(dtype) adjnuf_ob=adjnuf_ob.to(gpu) smapT=torch.ones((1,1,2,nx,nx)).cuda() nuf_ob = MriSenseNufft(im_size=(nx,nx),smap=smapT).to(dtype) nuf_ob=nuf_ob.to(gpu) adjnuf_ob = AdjMriSenseNufft(im_size=(nx,nx), smap=smapT).to(dtype) adjnuf_ob=adjnuf_ob.to(gpu)