<jupyter_start><jupyter_text>Setup Network<jupyter_code>Phi = TVM_Net(A.to(device))
Phi = Phi.to(device)
print(Phi)
pytorch_total_params = sum(p.numel() for p in Phi.parameters() if p.requires_grad)
print(f'Number of trainable parameters: {pytorch_total_params}')

max_epochs = 2000 
max_depth  = 250 
criterion  = torch.nn.MSELoss()  
fmt        = '[{:2d}/{:2d}]: train_loss = {:7.3e} | ' 
fmt       += 'depth = {:5.1f} | lr = {:5.1e} | time = {:4.1f} sec'<jupyter_output>TVM_Net(
  (shrink): Softshrink(0.1)
)
Number of trainable parameters: 0<jupyter_text>Train the Network<jupyter_code>best_loss = 1.0e10 

for epoch in range(max_epochs): 
  sleep(0.5)  # slows progress bar so it won't print on multiple lines
  tot = len(data_loader)
  loss_ave = 0.0
  start_time_epoch = time.time() 
  with tqdm(total=tot, unit=" batch", leave=False, ascii=True) as tepoch:      
    for idx, (u_batch, d) in enumerate(data_loader): 
        u_batch    = u_batch.to(device) 
        batch_size = u_batch.shape[0]
        train_batch_size = d.shape[0] # re-define if batch size changes
        u = Phi(d.to(device), max_depth=max_depth) # add snippet for hiding
        output = criterion(u, u_batch)
        train_loss = output.detach().cpu().numpy()
        loss_ave += train_loss * train_batch_size
 
        if idx%2 == 0:
            # compute test image 
            Phi.eval()
            u_test_approx = Phi(data_obs_test[0:1,:,:,:], max_depth=max_depth)

            plt.figure()
            plt.subplot(2,2,1)
            plt.imshow(u_batch[0,0,:,:].cpu(), vmin=0, vmax=1)
            plt.title('u true train')
            plt.subplot(2,2,2)
            plt.imshow(u[0,0,:,:].detach().cpu(), vmin=0, vmax=1)
            plt.title('u approx train')
            plt.subplot(2,2,3)
            plt.imshow(u_test[0,0,:,:].cpu(), vmin=0, vmax=1)
            plt.title('u true test')
            plt.subplot(2,2,4)
            plt.imshow(u_test_approx[0,0,:,:].detach().cpu(), vmin=0, vmax=1)
            plt.title('u approx test')
            plt.show() 
            Phi.train() 

        tepoch.update(1)
        tepoch.set_postfix(train_loss="{:5.2e}".format(train_loss),
                            depth="{:5.1f}".format(Phi.depth))
    

  loss_ave = loss_ave/len(data_loader.dataset)
  end_time_epoch = time.time()
  time_epoch = end_time_epoch - start_time_epoch
  #lr_scheduler.step()
  print(fmt.format(epoch+1, max_epochs, loss_ave, Phi.depth, 
                   0.0,
                   time_epoch))
from skimage.metrics import structural_similarity as ssim
from skimage.metrics import peak_signal_noise_ratio as psnr

n_samples = u_test.shape[0]

data_test     = TensorDataset(u_test, data_obs_test)
test_data_loader     = DataLoader(dataset=data_test, batch_size=batch_size, shuffle=True)


def compute_avg_SSIM_PSNR(u_true, u_gen, n_mesh, data_range):
    # assumes images are size n_samples x n_features**2 and are detached
    n_samples = u_true.shape[0]
    u_true = u_true.reshape(n_samples, n_mesh, n_mesh).cpu().numpy()
    u_gen  = u_gen.reshape(n_samples, n_mesh, n_mesh).cpu().numpy()
    ssim_val = 0
    psnr_val = 0
    for j in range(n_samples):
        ssim_val = ssim_val + ssim(u_true[j,:,:], u_gen[j,:,:], data_range=data_range)
        psnr_val = psnr_val + psnr(u_true[j,:,:], u_gen[j,:,:], data_range=data_range)
    return ssim_val/n_samples, psnr_val/n_samples

test_loss_ave = 0
test_PSNR_ave = 0
test_SSIM_ave = 0
with torch.no_grad():
    for idx, (u_batch, d) in enumerate(test_data_loader): 

        u_batch    = u_batch.to(device) 
        batch_size = u_batch.shape[0]
        temp       = u_batch.view(batch_size, -1)
        temp       = temp.permute(1,0)        
        test_batch_size = d.shape[0] # re-define if batch size changes
        Phi.eval()
        u = Phi(d, max_depth=max_depth) # add snippet for hiding
        output = criterion(u, u_batch)
        test_loss = output.detach().cpu().numpy()
        test_SSIM, test_PSNR = compute_avg_SSIM_PSNR(u_batch, u, 128, 1)
        test_PSNR_ave += test_PSNR * test_batch_size
        test_loss_ave += test_loss * test_batch_size
        test_SSIM_ave += test_SSIM * test_batch_size

        print('test_PSNR = {:7.3e}'.format(test_PSNR))
        print('test_SSIM = {:7.3e}'.format(test_SSIM))
        print('test_loss = {:7.3e}'.format(test_loss))
        if idx%1 == 0:
            # compute test image 
            plt.figure()
            plt.subplot(1,2,1)
            plt.imshow(u_batch[0,0,:,:].cpu(), vmin=0, vmax=1)
            plt.title('u true')
            plt.subplot(1,2,2)
            plt.imshow(u[0,0,:,:].detach().cpu(), vmin=0, vmax=1)
            plt.title('u approx')
            plt.show()  

print('\n\nSUMMARY')
print('test_loss_ave =  {:7.3e}'.format(test_loss_ave / len(data_loader.dataset)))
print('test_PSNR_ave =  {:7.3e}'.format(test_PSNR_ave / len(data_loader.dataset)))
print('test_SSIM_ave =  {:7.3e}'.format(test_SSIM_ave / len(data_loader.dataset)))
if use_ellipses:
    ind_val = 0
else:
    ind_val = 1000

u = Phi(data_obs_test[ind_val,:,:,:], max_depth=max_depth).view(128,128)
u_true = u_test[ind_val,0,:,:]
def string_ind(index):
    if index < 10:
        return '000' + str(index)
    elif index < 100:
        return '00' + str(index)
    elif index < 1000:
        return '0' + str(index)
    else:
        return str(index)

cmap = 'gray'
fig = plt.figure()
plt.imshow(np.rot90(u.detach().cpu().numpy()),cmap=cmap, vmin=0, vmax=1)
plt.axis('off')

data_type = 'Ellipse' if use_ellipses else 'Lodopab'

save_loc = './drive/MyDrive/FixedPointNetworks/Learned_Feasibility_' + data_type + '_TVM_ind_' + string_ind(ind_val) + '.pdf'
plt.savefig(save_loc,bbox_inches='tight')
plt.show()

print("SSIM: ", compute_avg_SSIM_PSNR(u_true.view(1,128,128), u.view(1,128,128).detach(), 128, 1))

############
# TRUE
###########

cmap = 'gray'
fig = plt.figure()
plt.imshow(np.rot90(u_true.detach().cpu().numpy()),cmap=cmap, vmin=0, vmax=1)
plt.axis('off')


save_loc = './drive/MyDrive/FixedPointNetworks/Learned_Feasibility_' + data_type + '_GT_ind_' + string_ind(ind_val) + '.pdf'
plt.savefig(save_loc,bbox_inches='tight')
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15337423312883436, 'fraction_numerical': 0.027123022279625444, 'mean_word_length': 3.0837178642056693, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '101232', 'license': 'MIT', 'n_tokens_mistral': 2539, 'n_tokens_neox': 2427, 'n_words': 406, 'path': 'TVM_Ellipses_LoDoPab.ipynb', 'repo_name': 'howardheaton/feasibility_fixed_point_networks', 'size': '6194'}