<jupyter_start><jupyter_text>Testing High-Dimensional Benchmarks - MM-1SolverDenseICNN-based benchmark training.**GPU-only implementation.**<jupyter_code>import os, sys
sys.path.append("..")

import matplotlib
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline 

import numpy as np
import torch
import gc

from sklearn.decomposition import PCA

from src import distributions
from src import map_benchmark

from src.icnn import DenseICNN
from src.tools import train_identity_map, unfreeze, freeze
from src.plotters import plot_benchmark_emb, plot_W2, plot_benchmark_metrics
import src.map_benchmark as mbm
from src.metrics import score_fitted_maps, score_baseline_maps, metrics_to_dict

import pandas as pd

from tqdm import tqdm
from IPython.display import clear_output<jupyter_output><empty_output><jupyter_text>Changable Config<jupyter_code>DIM = 16
BATCH_SIZE = 1024

GPU_DEVICE = 1<jupyter_output><empty_output><jupyter_text>Hardcoded Config<jupyter_code>BENCHMARK = 'Mix3toMix10'

LR = 1e-3
LAMBDA = DIM
MAX_ITER = 20001
INV_MAX_ITER = 1000
INV_TOL=1e-3

D_HYPERPARAMS = {
    'dim' : DIM,
    'rank' : 1,
    'hidden_layer_sizes' : [max(2*DIM, 64), max(2*DIM, 64), max(DIM, 32)],
    'strong_convexity' : 1e-4
}

OUTPUT_PATH = '../logs/' + BENCHMARK
if not os.path.exists(OUTPUT_PATH):
    os.makedirs(OUTPUT_PATH)

METHOD = 'MMv1'

assert torch.cuda.is_available()
torch.cuda.set_device(GPU_DEVICE)
torch.manual_seed(0x000000); np.random.seed(0x000000)<jupyter_output><empty_output><jupyter_text>Benchmark Setup<jupyter_code>benchmark = mbm.Mix3ToMix10Benchmark(DIM)
emb_X = PCA(n_components=2).fit(benchmark.input_sampler.sample(2**14).cpu().detach().numpy())
emb_Y = PCA(n_components=2).fit(benchmark.output_sampler.sample(2**14).cpu().detach().numpy())<jupyter_output><empty_output><jupyter_text>Discriminators Setup<jupyter_code>D = DenseICNN(**D_HYPERPARAMS).cuda()
D_conj = DenseICNN(**D_HYPERPARAMS).cuda()

pretrain_sampler = distributions.StandardNormalSampler(dim=DIM)
print('Pretraining identity potential. Final MSE:', train_identity_map(D, pretrain_sampler, convex=True, blow=3))
D_conj.load_state_dict(D.state_dict())
del pretrain_sampler<jupyter_output>Pretraining identity potential. Final MSE: 0.0008262658957391977<jupyter_text>Main training<jupyter_code>D_opt = torch.optim.Adam(D.parameters(), lr=LR)
D_conj_opt = torch.optim.Adam(D_conj.parameters(), lr=LR)
W2_history = []
metrics = dict(L2_UVP_fwd=[], cos_fwd=[], L2_UVP_inv=[], cos_inv=[])
baselines = {
    baseline : metrics_to_dict(*score_baseline_maps(benchmark, baseline))
    for baseline in ['identity', 'constant', 'linear']
}
L2_UVP_fwd_min, L2_UVP_inv_min = np.inf, np.inf
def invert(D, Y, lr=0.1, max_iter=INV_MAX_ITER, tol=INV_TOL):
    freeze(D)
    Y_inv = torch.rand_like(Y, requires_grad=True)

    opt = torch.optim.Adam([Y_inv], lr=lr)
    torch.sum(Y_inv).backward() # Terrible solution to init gradients
    opt.zero_grad()

    for it in range(max_iter):
        loss = D(Y_inv).sum() - (Y * Y_inv).sum()
        loss.backward()

        if torch.sqrt((Y_inv.grad.data ** 2).mean(dim=1)).mean() < tol:
            break

        opt.step()
        opt.zero_grad()
    return Y_inv.detach()
print('Fitting the monotone generative mapping')
for iteration in tqdm(range(MAX_ITER)):
    X = benchmark.input_sampler.sample(BATCH_SIZE); X.requires_grad_(True)
    Y = benchmark.output_sampler.sample(BATCH_SIZE); Y.requires_grad_(True)
        
    Y_inv = invert(D, Y)
    unfreeze(D)
    D_opt.zero_grad(); 
    W_loss_XY = (D(X) - D(Y_inv)).mean()
    W_loss_XY.backward()
    D_opt.step(); D.convexify();
    
    with torch.no_grad():
        W_loss_nograd_XY = (- (X ** 2).sum(dim=1) / 2).mean() +\
        ((Y_inv * Y).sum(dim=1) - (Y_inv ** 2).sum(dim=1) / 2 ).mean()
    
    X_inv = invert(D_conj, X)
    unfreeze(D_conj)
    D_conj_opt.zero_grad()
    W_loss_YX = (D_conj(Y) - D_conj(X_inv)).mean()
    W_loss_YX.backward()
    D_conj_opt.step(); D_conj.convexify()
    
    with torch.no_grad():
        W_loss_nograd_YX = (- (Y ** 2).sum(dim=1) / 2).mean() +\
        ((X_inv * X).sum(dim=1) - (X_inv ** 2).sum(dim=1) / 2 ).mean()

    W_loss = (W_loss_XY + W_loss_YX) / 2
    with torch.no_grad():
        W_loss_nograd = (W_loss_nograd_XY + W_loss_nograd_YX) / 2
     
    W2_history.append(-W_loss.item() - W_loss_nograd.item())
    
    if iteration % 100 == 0:
        L2_UVP_fwd, cos_fwd, L2_UVP_inv, cos_inv = score_fitted_maps(benchmark, D, D_conj)        
        metrics['L2_UVP_fwd'].append(L2_UVP_fwd); metrics['cos_fwd'].append(cos_fwd)
        metrics['L2_UVP_inv'].append(L2_UVP_inv); metrics['cos_inv'].append(cos_inv)
               
        if L2_UVP_fwd < L2_UVP_fwd_min:
            torch.save(D.state_dict(), os.path.join(OUTPUT_PATH, f'{DIM}-{METHOD}.pt'))
            L2_UVP_fwd_min = L2_UVP_fwd
    
        if L2_UVP_inv < L2_UVP_inv_min:
            torch.save(D_conj.state_dict(), os.path.join(OUTPUT_PATH, f'{DIM}-{METHOD}_conj.pt'))
            L2_UVP_inv_min = L2_UVP_inv
            
        results = pd.DataFrame(metrics)
        results.to_csv(os.path.join(OUTPUT_PATH, f'{DIM}-{METHOD}.csv'))
                    
    if iteration % 1000 == 0:
        clear_output(wait=True)
        print("Iteration", iteration)
        
        fig, axes = plot_benchmark_emb(benchmark, emb_X, emb_Y, D, D_conj)
        plt.show(fig); plt.close(fig)
        
        fig, axes = plot_W2(benchmark, W2_history)
        plt.show(fig); plt.close(fig)
        
        fig, axes = plot_benchmark_metrics(benchmark, metrics, baselines)
        plt.show(fig); plt.close(fig)<jupyter_output>Iteration 20000
{'fraction_non_alphanumeric': 0.13580028328611898, 'fraction_numerical': 0.024610481586402267, 'mean_word_length': 3.9079061685490877, 'pattern_counts': {'":': 0, '<': 26, '<?xml version=': 0, '>': 23, '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': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '520905', 'license': 'MIT', 'n_tokens_mistral': 2341, 'n_tokens_neox': 2229, 'n_words': 356, 'path': 'notebooks/MMv1_test_hd_benchmark.ipynb', 'repo_name': 'iamalexkorotin/Wasserstein2Benchmark', 'size': '5648'}