<jupyter_start><jupyter_text>only use first and last point<jupyter_code>from weight_norm import weight_norm
set_random_seeds(20190118, False)
discriminator = nn.Sequential(nn.Linear(1,60), 
                             nn.ReLU(),#UpperBoundReLU(20)
                             nn.Linear(60,60),
                             nn.ReLU(),#UpperBoundReLU(20)
                             nn.Linear(60,1))


weight_norms = []
for m in discriminator.modules():
    if hasattr(m, 'weight') and m.weight is not None:
        w_norms = th.ones_like(m.weight[:,0:1], requires_grad=True)
        w_norms.data = w_norms.data *  1
        # in-place!!
        print("wnorm", w_norms.shape)
        weight_norm(m, 'weight', dim=0, fixed_norm=w_norms)
        weight_norms.append(w_norms)
None
optimizer = th.optim.Adam(discriminator.parameters(), lr=1e-2)
optim_wnorms = th.optim.Adam(weight_norms, lr=1e-3)
n_epochs = 20
for i_epoch in range(n_epochs):
    used_points = example_points[[0,2]]
    out = th.mean(discriminator(used_points)) - th.mean(discriminator(counter_points))
    l1 = used_points.unsqueeze(1)- counter_points.unsqueeze(0)
    dist = discriminator(used_points).unsqueeze(1) - discriminator(counter_points).unsqueeze(0)
    dist_loss = nn.functional.relu(th.abs(dist) - th.abs(l1))
    dist_loss =  dist_loss + (dist_loss ** 2 )
    loss = -out + th.mean(dist_loss) * 10
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    used_points = example_points[[1,]]
    out = th.mean(discriminator(used_points)) - th.mean(discriminator(counter_points))
    l1 = used_points.unsqueeze(1)- counter_points.unsqueeze(0)
    dist = discriminator(used_points).unsqueeze(1) - discriminator(counter_points).unsqueeze(0)
    dist_loss = nn.functional.relu(th.abs(dist) - th.abs(l1))
    dist_loss =  dist_loss + (dist_loss ** 2 )
    loss = -out + th.mean(dist_loss) * 10
    optim_wnorms.zero_grad()
    loss.backward()
    optim_wnorms.step()
    
    if i_epoch % (n_epochs // 20) == 0:
        print("loss", loss)
        sorted_p, _ = th.sort(th.cat((example_points, counter_points)), dim=0)

        outs = discriminator(sorted_p)
        fig = plt.figure(figsize=(8,3))
        plt.plot(var_to_np(sorted_p.squeeze()), var_to_np(outs))
        display(fig)
        plt.close(fig)
from weight_norm import weight_norm
set_random_seeds(20190118, False)
discriminator = nn.Sequential(nn.Linear(1,60), 
                             nn.ReLU(),#UpperBoundReLU(20)
                             nn.Linear(60,60),
                             nn.ReLU(),#UpperBoundReLU(20)
                             nn.Linear(60,1))


weight_norms = []
for m in discriminator.modules():
    if hasattr(m, 'weight') and m.weight is not None:
        w_norms = th.ones_like(m.weight[:,0:1], requires_grad=True)
        w_norms.data = w_norms.data *  1
        # in-place!!
        print("wnorm", w_norms.shape)
        weight_norm(m, 'weight', dim=0, fixed_norm=w_norms)
        weight_norms.append(w_norms)
None
a = th.randn(784,2000)
%%timeit
_, _ = th.sort(a, dim=0)
optimizer = th.optim.Adam(discriminator.parameters(), lr=1e-2)
optim_wnorms = th.optim.Adam(weight_norms, lr=1e-3)
n_epochs = 2000
for i_epoch in range(n_epochs):
    used_points = example_points[th.randperm(len(example_points))[:2]]
    out = th.mean(discriminator(used_points)) - th.mean(discriminator(counter_points))
    l1 = used_points.unsqueeze(1)- counter_points.unsqueeze(0)
    dist = discriminator(used_points).unsqueeze(1) - discriminator(counter_points).unsqueeze(0)
    dist_loss = nn.functional.relu(th.abs(dist) - th.abs(l1))
    dist_loss =  dist_loss + (dist_loss ** 2 )
    loss = -out + th.mean(dist_loss) * 10
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    used_points = example_points[th.randperm(len(example_points))[:1]]
    out = th.mean(discriminator(used_points)) - th.mean(discriminator(counter_points))
    l1 = used_points.unsqueeze(1)- counter_points.unsqueeze(0)
    dist = discriminator(used_points).unsqueeze(1) - discriminator(counter_points).unsqueeze(0)
    dist_loss = nn.functional.relu(th.abs(dist) - th.abs(l1))
    dist_loss =  dist_loss + (dist_loss ** 2 )
    loss = -out + th.mean(dist_loss) * 10
    optim_wnorms.zero_grad()
    loss.backward()
    optim_wnorms.step()
    
    if i_epoch % (n_epochs // 20) == 0:
        print("loss", loss)
        sorted_p, _ = th.sort(th.cat((example_points, counter_points)), dim=0)

        outs = discriminator(sorted_p)
        fig = plt.figure(figsize=(8,3))
        plt.plot(var_to_np(sorted_p.squeeze()), var_to_np(outs))
        display(fig)
        plt.close(fig)
import ot
used_points.shape
distances = np.sum(np.abs(var_to_np(used_points)[:,None] - var_to_np(counter_points)[None]), axis=(2))
coupling = ot.emd([],[], distances)
np.sum(coupling * distances)
distac<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1403688524590164, 'fraction_numerical': 0.02622950819672131, 'mean_word_length': 3.1969045571797077, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, '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': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '603240', 'license': 'MIT', 'n_tokens_mistral': 1858, 'n_tokens_neox': 1794, 'n_words': 265, 'path': 'notebooks/BhNo_from_2_to_64_samples/1d_Regulate_Norm_Idea.ipynb', 'repo_name': 'robintibor/reversible2', 'size': '4880'}