#def for_nufft(nuff,x,csm,traj):
#def for_nufft(nufft_ob,x,csm,traj):
#    nch=csm.shape[1]
#    nbas=x.shape[0]
#    tmp1=torch.zeros((nch,nbas,2,NF*nx*nintl)).cuda()
#    for i in range(nch):
#        tmp=sf.complex_mult(x,csm[:,i].repeat(nbas,1,1,1,1),dim=2)
#        tmp=torch.reshape(tmp,(nbas,1,2,nx,nx))
#        #xx=ktrajT[0].unsqueeze(0)
#        tmp1[i]=nufft_ob(tmp,traj.repeat(nbas,1,1).cuda()).squeeze(1)
#    tmp1=tmp1.permute(1,0,2,3)
#    return tmp1

#%%
G = UnetClass().to(gpu)
G = SmallModel1().to(gpu)

GV = SmallModel().to(gpu)
#G.load_state_dict(torch.load('wts-10U2.pt'))
#GV.load_state_dict(torch.load('wts-10V2.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-3}])

#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,
#recT=torch.load('stm900U.pt')
atbT = torch.load('atb200.pt')

#%% take them to gpu
#B=B.to(gpu)
#atbT=atbT.to(gpu)

VT = VT.to(gpu)
recT = recT.to(gpu)

#%% Loading trained networks parameters

wts1 = dir2 + 'wts-20.pt'
wts2 = dir3 + 'wts-500.pt'

gnu = SmallModel1().to(gpu)
gnu.load_state_dict(torch.load(wts1))

gnv = SmallModel().to(gpu)
gnv.load_state_dict(torch.load(wts2))

wtsFname1 = 'wts-10U' + str(2) + '.pt'
#wtsFname1='wts-2U'+str(60+1)+'.pt'

G = UnetClass().to(gpu)
G.load_state_dict(torch.load(wtsFname1))

#atbT=atbT*W
z = recT.permute(3, 0, 1, 2)
z = torch.reshape(z, (1, 2 * nbasis, 512, 512)).to(gpu)
#z = torch.randn((1,60,512,512),device=gpu, dtype=dtype)
예제 #3
0
def myPSNR(org,recon):
    sqrError=abs(org-recon)**2
    N=torch.prod(torch.tensor(org.shape[-2:]))
    mse=sqrError.sum(dim=(-1,-2))
    mse=mse/N
    
    #maxval=np.max(org,axis=(-1,-2)) + 1e-15
    psnr=10*torch.log10(mx**2/(mse+1e-15))

    return psnr
        
#%%        
        #G=UnetClass().to(gpu)
        G=SmallModel1().to(gpu)
        G.load_state_dict(torch.load('tempp_10000.pt'))
        optimizer=torch.optim.Adam([{'params':G.parameters(),'lr':1e-4}])
        scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.7, patience=6, verbose=True, min_lr=1e-5)
        #z1=recT+torch.normal(0,0.02,(recT1.shape[0],recT1.shape[1],recT1.shape[2],recT1.shape[3],recT1.shape[4])).cuda()
        z=recT1
        z=torch.reshape(z,(1,2*nbasis,nx,nx)).to(gpu)
        #z1=torch.cuda.FloatTensor(3,2*nbasis,nx,nx).fill_(0) 
        #noi=[0.0003,0.05,0.01]
        #noi=[0.003,0.0005,0.01]
        #noi=[0.05,0.0005,0.01]
        #noi=[0.05,0.01,0.1]

#        for it in range(3):
#            z1[it]=z+torch.normal(0,noi[it],(z.shape[0],z.shape[1],z.shape[2],z.shape[3])).cuda()
#        z=z.repeat(3,1,1,1)
        pp=np.array([])

        sf.tic()
        for rn in range(40):
            #z=atb#+torch.normal(0,noi[it],(atb.shape[0],atb.shape[1],atb.shape[2],atb.shape[3])).cuda()
            #z=torch.reshape(z,(1,2*nbasis,nx,nx))
                #z=atb+torch.normal(0,0.005,(atb.shape[0],atb.shape[1],atb.shape[2],atb.shape[3])).cuda()
                #z=torch.reshape(z,(1,2*nbasis,nx,nx))
            u1=G(z)#+z.cuda()
            u1=torch.reshape(u1,(nbasis,1,2,nx,nx))
            l1_reg=0
            for param in G.parameters():
                l1_reg += param.abs().sum()
            loss=((abs(u1-recT1))).pow(2).sum()+1.0*l1_reg
            if rn%20==0:
#                print(rn,loss.item())
            #torch.cuda.empty_cache()
                plt.figure(rn)
                dd=(v2[0,0,:,20].unsqueeze(1).unsqueeze(1)*u1.squeeze(1)).sum(dim=0).detach().cpu().numpy()
                plt.imshow(np.abs((dd[0,100:400,100:400]+dd[1,100:400,100:400]*1j)),cmap='gray')
                plt.show()
                plt.pause(1)
            
            dd=(v2[0,0,:,10:50].unsqueeze(-1).unsqueeze(-1).detach().cpu()*u1[:,:,:,100:400,100:400].detach().cpu()).sum(dim=0).detach().cpu()
            org1=(v2[0,0,:,10:50].unsqueeze(-1).unsqueeze(-1).detach().cpu()*recT[:,:,:,100:400,100:400].detach().cpu()).sum(dim=0)
            psnr=torch.mean(myPSNR(org1,dd))
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            scheduler.step(loss.detach())
            print(rn,loss.item(),psnr.item())
            pp=np.append(pp,psnr)
            
        sf.toc()
        
        #%%
        #        sf.tic()
        #z1=torch.normal(0,0.01,(1,1,2,nf1)).cuda()
        z1=0.01*torch.ones((1,1,2,nf1)).to(gpu)
        z1[0,0,1,:]=torch.tensor(V[22,0:nf1])
        z1[0,0,0,:]=torch.tensor(V[27,0:nf1])

        #vv=V[30:32]
        #z1=torch.reshape(torch.tensor(V[29:31,0:nf1]),(1,1,2,nf1)).to(gpu)
        z1 = Variable(z1,requires_grad=True)
        V1=torch.tensor(V[:,0:nf1]).cuda()#+torch.normal(0,0.01,(V.shape[0],V.shape[1])).cuda()
        #V1=torch.reshape(V1,(1,1,nbasis,nf1))
        GV=SmallModel(16).to(gpu)
        #GV=generatorV(2).to(gpu)
        optimizer=torch.optim.Adam([{'params':GV.parameters(),'lr':1e-4},{'params':z1,'lr':1e-6}])
        scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.7, patience=6, verbose=True, min_lr=1e-7)
        pp=np.array([])
        for ep1 in range(25000):
            for bat in range(1):
                
                #z1=torch.reshape(V1,(1,1,nbasis,nf1))
                l1_reg=0.
                for param in GV.parameters():
                    l1_reg += param.abs().sum()  
                v1=GV(z1)
                #v1=v1.permute(1,0)
                loss=abs(v1[0,0]-V1).pow(2).sum()+0.000*l1_reg+10.0*Smoothness(z1[0,0,:,:]) 
                if ep1%10==0:
                    print(ep1,loss.item())
                
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                scheduler.step(loss.detach())
                #pp=np.append(pp,loss.item())

        
        del kdaT,smap,adjnufft_ob
        #%%
        def Smoothness(zvec):        
            zsmoothness = zvec[:,1:]-zvec[:,:-1]
            zsmoothness = torch.sum(zsmoothness*zsmoothness,axis=1).squeeze()
            zsmoothness = torch.sum(zsmoothness,axis=0)
            return(zsmoothness)
        
        z=recT1#+torch.normal(0,0.03,(atb.shape[0],atb.shape[1],atb.shape[2],atb.shape[3])).cuda()#.permute(1,0,2,3)
        z=torch.reshape(z,(1,2*nbasis,nx,nx))
        #z1=torch.cuda.FloatTensor(1,1,nbasis,nf1).fill_(0)
        #z1=torch.normal(0,1,(1,1,nbasis,nf1))
        z1=torch.reshape(torch.tensor(V[:,0:nf1]),(1,1,nbasis,nf1))
        #z1 = Variable(z1,requires_grad=True)
        #%%    
        torch.cuda.empty_cache()
        #from dn_modelU2 import SmallModel1
#        G=SmallModel1().to(gpu)
#        GV=SmallModel().to(gpu)
#        G.load_state_dict(torch.load('tempp_199.pt'))
##        G.load_state_dict(torch.load('wtsN-U3.pt'))
##        GV.load_state_dict(torch.load('wtsN-V3.pt'))        
#        #optimizer=torch.optim.AdamW([{'params':G.parameters(),'lr':1e-4},{'params':GV.parameters(),'lr':1e-4},{'params':z,'lr':1e-4}])
        #v1=(z1.cuda()).unsqueeze(-1)
        #v1 = Variable(v1,requires_grad=True)
        optimizer=torch.optim.Adam([{'params':G.parameters(),'lr':1e-5},{'params':GV.parameters(),'lr':1e-4},{'params':z1,'lr':1e-6}])
#        #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)
        
        #optimizer=torch.optim.SGD([{'params':v1,'lr':1e-4}])#,{'params':z,'lr':5e-3,'momentum':0.9}])
        
        smapT=smapT.cpu()
        ktrajT=ktrajT.to(cpu)
        dcf=torch.reshape(dcf,(nbasis,3,2,NF,nintl*nx))
        plt.figure(2)
        dd=(np.reshape(V[:,26],(nbasis,1,1,1))*recT.detach().cpu().numpy()).sum(axis=0)
        plt.imshow(np.abs((dd[0,100:400,100:400]+dd[1,100:400,100:400]*1j)),cmap='gray')
        pp=np.array([])
        plt.pause(1)
        plt.figure(1)
        sf.tic()
        for it in range(500):    
#            for bat in range(1):
                
            v1=GV(z1.cuda()).unsqueeze(-1)
            u1=G(z.cuda())#+z.cuda()
            u1=torch.reshape(u1,(nbasis,2,nx,nx))
            u1=u1.unsqueeze(1)
                  
            l1_regU=0.
            for param in G.parameters():
                l1_regU += param.abs().sum()
                
            l1_regV=0.
            for param in GV.parameters():
                l1_regV += param.abs().sum()
                
            err=abs(AUV1(u1,v1)-kdataT1[:,:,0:nf1].cuda()).pow(2)
            loss=((dcf[0,:,:,0:nf1])*(err)).sum()+0.000005*l1_regU+0.00001*l1_regV+10.0*Smoothness(z1[0,0,:,:]) 
#            err=abs(yy1.squeeze(1)-atb).pow(2)          
#            loss=(err).sum()#+10*l1_reg#+0.1*(z.cpu()-u2.cpu()).pow(2).sum()+0.1*(z1[0,0].cpu()-v2.cpu()).pow(2).sum()
            if it%10==0:
                plt.figure(it)
                dd=(v1[0,0,:,20].unsqueeze(1).unsqueeze(1)*u1.squeeze(1)).sum(dim=0).detach().cpu().numpy()
                plt.imshow(np.abs((dd[0,100:400,100:400]+dd[1,100:400,100:400]*1j)),cmap='gray')
                plt.show()
                plt.pause(1)           
#            if it%10==0:
#                torch.save(u1,'tmpU'+str(it)+'.pt')           
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            scheduler.step(loss.detach())
            
            dd=(v2[0,0,:,10:50].unsqueeze(-1).unsqueeze(-1).detach().cpu()*u1[:,:,:,100:400,100:400].detach().cpu()).sum(dim=0).detach().cpu()
            org1=(v2[0,0,:,10:50].unsqueeze(-1).unsqueeze(-1).detach().cpu()*recT[:,:,:,100:400,100:400].detach().cpu()).sum(dim=0)
            psnr=torch.mean(myPSNR(org1,dd))
            print(it,loss.item(),psnr.item())
            pp=np.append(pp,psnr)
        sf.toc()
예제 #4
0
n_select = 30
nbasis = 30
#%%Generate a meaningful filename to save the trainined models for testing
print('*************************************************')
start_time = time.time()
saveDir = 'UPTmodels/'
cwd = os.getcwd()
directory=saveDir+datetime.now().strftime("%d%b_%I%M%S%P_")+ \
 str(epochs)+'ep_' +'27oct'

if not os.path.exists(directory):
    os.makedirs(directory)
sessFileName = directory + '/modelUb'

#%% creating the training model
unet = SmallModel1()
unet = unet.to(gpu)
optimizer = torch.optim.Adam(unet.parameters(), lr=1e-4)


def lossFun(pred, org):
    loss = torch.mean(torch.abs(pred - org))
    return loss


#%% training code
#print ('training started on', datetime.now().strftime("%d-%b-%Y at %I:%M %P"))
#start_time=time.time()
torch.save(unet.state_dict(), directory + '/wts-0.pt')
#writer = SummaryWriter(log_dir=directory+'/')
#writer.add_graph(unet,torch.randn(1,2,512,512,30).to(gpu))