<jupyter_start><jupyter_text>set feature model weights and distribution to good start parameters<jupyter_code>from reversible2.constantmemory import clear_ctx_dicts
from reversible2.distribution import TwoClassDist

feature_model.data_init(th.cat((train_inputs[0], train_inputs[1]), dim=0))

# Check that forward + inverse is really identical
t_out = feature_model(train_inputs[0][:2])
inverted = invert(feature_model, t_out)
clear_ctx_dicts(feature_model)
assert th.allclose(train_inputs[0][:2], inverted, rtol=1e-3,atol=1e-4)
device = list(feature_model.parameters())[0].device
from reversible2.ot_exact import ot_euclidean_loss_for_samples
class_dist = TwoClassDist(2, np.prod(train_inputs[0].size()[1:]) - 2)
class_dist.cuda()

for i_class in range(2):
    with th.no_grad():
        this_outs = feature_model(train_inputs[i_class])
        mean = th.mean(this_outs, dim=0)
        std = th.std(this_outs, dim=0)
        class_dist.set_mean_std(i_class, mean, std)
        # Just check
        setted_mean, setted_std = class_dist.get_mean_std(i_class)
        assert th.allclose(mean, setted_mean)
        assert th.allclose(std, setted_std)
clear_ctx_dicts(feature_model)

optim_model = th.optim.Adam(feature_model.parameters(), lr=1e-4, betas=(0,0.9))
optim_dist = th.optim.Adam(class_dist.parameters(), lr=1e-2, betas=(0,0.9))

class ResidualBlock(nn.Module):
    def __init__(self, F):
        super(ResidualBlock, self).__init__()
        self.F = F

    def forward(self, x):
        return self.F(x) + x


from reversible2.discriminator import ProjectionDiscriminator
def res_block(n_c, n_i_c):
     return ResidualBlock(
        nn.Sequential(
        nn.Conv2d(n_c, n_i_c, (3,1), stride=1, padding=(1,0),bias=True),
        nn.ReLU(),
            nn.Conv2d(n_i_c, n_c, (3,1), stride=1, padding=(1,0),bias=True)),
    )
    


from braindecode.torch_ext.util import set_random_seeds
adv_model = nn.Sequential(
    nn.Conv2d(22,64, (3,1), stride=1, padding=(1,0),bias=True), #256
    res_block(64,96),
    res_block(64,96),
    nn.AvgPool2d((2,1)), #128
    res_block(64,96),
    res_block(64,96),
    nn.AvgPool2d((2,1)), # 64
    res_block(64,96),
    res_block(64,96),
    nn.AvgPool2d((2,1)), # 32
    res_block(64,96),
    res_block(64,96),
    nn.AvgPool2d((2,1)), # 16
    res_block(64,96),
    res_block(64,96),
    nn.AvgPool2d((2,1)), # 8
    res_block(64,96),
    res_block(64,96),
    nn.AvgPool2d((2,1)), # 4
    res_block(64,96),
    res_block(64,96),
    nn.AvgPool2d((2,1)), # 2
    res_block(64,96),
    res_block(64,96),
    ViewAs((-1,64,2,1), (-1,128)),
    )

adv_model = ProjectionDiscriminator(adv_model,128,2,)
adv_model.cuda();

optim_adv = th.optim.Adam([{
    'params': adv_model.parameters(),
    'lr': 4e-4, 'weight_decay': 0.00}],#lr 0.0004
                         betas=(0,0.9))
from reversible2.graph import data_init_sequential_module
_ = data_init_sequential_module(adv_model.adv_feature_model, th.cat((train_inputs[0], train_inputs[1]), dim=0))
import pandas as pd
df = pd.DataFrame()
g_loss = np_to_var([np.nan],dtype=np.float32)
g_grad = np.nan
g_grad_norm = np.nan
d_loss = np_to_var([np.nan],dtype=np.float32)
gradient_loss = np_to_var([np.nan],dtype=np.float32)
d_grad = np.nan
d_grad_norm = np.nan
from reversible2.constantmemory import clear_ctx_dicts
from reversible2.timer import Timer
from plot import plot_outs
from reversible2.gradient_penalty import gradient_penalty


i_start_epoch_out = 401
n_epochs = 5001
gen_frequency = 10
for i_epoch in range(n_epochs):
    epoch_row = {}
    with Timer(name='EpochLoop', verbose=False) as loop_time:
        gen_update = (i_epoch % gen_frequency) == (gen_frequency-1)
            
        optim_model.zero_grad()
        optim_dist.zero_grad()
        optim_adv.zero_grad()
        for i_class in range(len(train_inputs)):
            y = np_to_var([i_class]).cuda()
            class_ins = train_inputs[i_class]
            other_class_ins = train_inputs[1-i_class]
            with th.set_grad_enabled(gen_update):
                samples = class_dist.get_samples(i_class, len(train_inputs[i_class]) * 4)
                inverted = feature_model.invert(samples)
            
            with th.set_grad_enabled(gen_update):
                outs = feature_model(class_ins)
                changed_to_other_class = class_dist.change_to_other_class(outs, i_class_from=i_class,
                                                                          i_class_to=1-i_class)
                other_inverted = feature_model.invert(changed_to_other_class)
            if not gen_update:
                score_fake = adv_model(inverted.detach(), y)
                score_real = adv_model(class_ins, y)
                gradient_loss = gradient_penalty(adv_model, class_ins, inverted[:(len(class_ins))].detach(), y)
                d_loss = -score_real.mean() + score_fake.mean() + gradient_loss * 100
                d_loss.backward()
                
                score_real_other = adv_model(other_class_ins.detach(), 1-y)
                score_fake_other = adv_model(other_inverted.detach(), 1-y)
                gradient_loss_other = gradient_penalty(adv_model, other_class_ins, 
                                                 other_inverted[:(len(class_ins))].detach(), 1-y)
                d_loss_other =  -score_real_other.mean() + score_fake_other.mean() + gradient_loss_other * 100
                d_loss_other.backward()
                
                # Clip gradient
                d_grad_norm = np.mean([th.nn.utils.clip_grad_norm_([p],100) for p in adv_model.parameters()])
                d_grad = np.mean([th.sum(p.grad **2).item() for p in adv_model.parameters()])
                wd_d = -(-score_real.mean() + score_fake.mean()).item()
                epoch_row['wd_d_{:d}'.format(i_class)] = wd_d
            else:
                score_fake = adv_model(inverted, y)
                score_fake_other = adv_model(other_inverted, 1-y)
                g_loss = -th.mean(score_fake) - th.mean(score_fake_other)
                g_loss.backward()
                # Clip gradient
                g_grad_norm = np.mean([th.nn.utils.clip_grad_norm_([p],100) for p in feature_model.parameters()])
                g_grad = np.mean([th.sum(p.grad **2).item() for p in feature_model.parameters()])
            clear_ctx_dicts(feature_model)
                
            """
            assert False
            ot_loss_in = ot_euclidean_loss_for_samples(class_ins.view(class_ins.shape[0], -1),
                                                       inverted.view(inverted.shape[0], -1))
            del inverted
            with th.set_grad_enabled(gen_update):
                outs = feature_model(class_ins)
            if i_epoch < i_start_epoch_out:
                ot_loss_out = th.zeros(1, device=class_ins.device)
            else:
                ot_loss_out = ot_euclidean_loss_for_samples(outs[:,:2].squeeze(), samples[:,:2].squeeze())
            del samples
                
            other_class_ins = train_inputs[1-i_class]
            changed_to_other_class = class_dist.change_to_other_class(outs, i_class_from=i_class, i_class_to=1-i_class)
            other_inverted = feature_model.invert(changed_to_other_class)
            ot_transformed_in = ot_euclidean_loss_for_samples(other_class_ins.view(other_class_ins.shape[0], -1),
                                                              other_inverted.view(other_inverted.shape[0], -1))
            
            if i_epoch < i_start_epoch_out:
                ot_transformed_out = th.zeros(1, device=class_ins.device)
            else:
                other_samples = class_dist.get_samples(1-i_class, len(train_inputs[i_class]) * 2)
                ot_transformed_out = ot_euclidean_loss_for_samples(changed_to_other_class[:,:2].squeeze(),
                                                                   other_samples[:,:2].squeeze(),)
            loss =  ot_loss_out + ot_transformed_in + ot_transformed_out + ot_loss_in
            loss.backward()
            del outs, other_class_ins, other_inverted
            clear_ctx_dicts(feature_model)"""
        if not gen_update:
            optim_adv.step()
        else:
            optim_model.step()
            optim_dist.step()
    epoch_row.update({
    'd_loss': d_loss.item(),
    'g_loss': g_loss.item(),
    'grad_loss': gradient_loss.item(),
    'o_real': th.mean(score_real).item(),
    'o_fake': th.mean(score_fake).item(),
    'g_grad': g_grad,
    'g_grad_norm': g_grad_norm,
    'd_grad': d_grad,
    'd_grad_norm': d_grad_norm,
    'runtime': loop_time.elapsed_secs * 1000})
    if i_epoch % (n_epochs // 20) != 0:
        df = df.append(epoch_row, ignore_index=True)
        # otherwise add ot loss in
    else:
        for i_class in range(len(train_inputs)):
            with th.no_grad():
                samples = class_dist.get_samples(i_class, len(train_inputs[i_class]) * 4)
                inverted = feature_model.invert(samples)
                clear_ctx_dicts(feature_model)
                ot_loss_in = ot_euclidean_loss_for_samples(class_ins.view(class_ins.shape[0], -1),
                                                           inverted.view(inverted.shape[0], -1)[:(len(class_ins))])
                epoch_row['ot_loss_in_{:d}'.format(i_class)] = ot_loss_in.item()
        df = df.append(epoch_row, ignore_index=True)
        print("Epoch {:d} of {:d}".format(i_epoch, n_epochs))
        print("G Loss: {:.2E}".format(g_loss.item()))
        print("D Loss: {:.2E}".format(d_loss.item()))
        print("OT Loss In: {:.2E}".format(ot_loss_in.item()))
        #print("OT Loss Out: {:.2E}".format(ot_loss_out.item()))
        #print("Transformed OT Loss In: {:.2E}".format(ot_transformed_in.item()))
        #print("Transformed OT Loss Out: {:.2E}".format(ot_transformed_out.item()))
        print("Loop Time: {:.0f} ms".format(loop_time.elapsed_secs * 1000))
        display(df.iloc[-3:])
        plot_outs(feature_model, train_inputs, test_inputs,
                 class_dist)
        fig = plt.figure(figsize=(8,2))
        plt.plot(var_to_np(th.cat((th.exp(class_dist.class_log_stds),
                                 th.exp(class_dist.non_class_log_stds)))),
                marker='o')
        display_close(fig)
    

with th.no_grad():
    for n_samples_factor in [2,1]:
        for setname, set_inputs in (("Train", train_inputs), ("Test", test_inputs)):
            total_loss = 0
            total_spec_ot = 0
            for i_class in range(2):
                class_ins = set_inputs[i_class]
                samples = class_dist.get_samples(i_class, len(train_inputs[i_class]) * n_samples_factor)
                inverted = feature_model.invert(samples)
                loss = ot_euclidean_loss_for_samples(class_ins.view(class_ins.shape[0],-1),
                                                     inverted.view(inverted.shape[0],-1))
                total_loss += loss
                spec_ot = ot_euclidean_loss_for_samples(
                    th.rfft(class_ins.squeeze(), signal_ndim=1, normalized=True).view(class_ins.shape[0],-1),
                    th.rfft(inverted.squeeze(), signal_ndim=1, normalized=True).view(inverted.shape[0],-1))
                total_spec_ot += spec_ot
                clear_ctx_dicts(feature_model)
            print("{:d} Samples ". format(len(samples)) + setname + " Loss: {:.1f}".format(total_loss.item() / 2))
            print("{:d} Samples ". format(len(samples)) + setname + " SpecLoss: {:.1f}".format(spec_ot.item() / 2))

    # Show losses between sets
    for setname, set_inputs in (("Train", train_inputs),):
        total_loss = 0
        total_spec_ot = 0
        for i_class in range(2):
            class_ins = set_inputs[i_class]
            samples = class_dist.get_samples(i_class, len(train_inputs[i_class]) * 5) # take same number of samples
            inverted = test_inputs[i_class]
            loss = ot_euclidean_loss_for_samples(class_ins.view(class_ins.shape[0],-1),
                                                     inverted.view(inverted.shape[0],-1))
            total_loss += loss
            spec_ot = ot_euclidean_loss_for_samples(
                th.rfft(class_ins.squeeze(), signal_ndim=1, normalized=True).view(class_ins.shape[0],-1),
                th.rfft(inverted.squeeze(), signal_ndim=1, normalized=True).view(inverted.shape[0],-1))
            total_spec_ot += spec_ot
        print(setname + " Loss: {:.1f}".format(total_loss.item() / 2))
        print(setname + " SpecLoss: {:.1f}".format(spec_ot.item() / 2))

from reversible2.plot import plot_head_signals_tight

mean_bps_per_class = []
for i_class in range(len(train_inputs)):
    inverted = train_inputs[i_class]
    mean_bps_per_class.append(
        np.mean(np.abs(np.fft.rfft(var_to_np(inverted.squeeze()))), axis=0))

fig = plt.figure(figsize=(8,3))
plt.plot(np.mean(mean_bps_per_class, axis=1).T)
display_close(fig)

fig = plot_head_signals_tight(np.stack(mean_bps_per_class, axis=-1), sensor_names=sensor_names,
                             figsize=(20,12));
from reversible2.plot import plot_head_signals_tight

mean_bps_per_class = []
for i_class in range(len(train_inputs)):
    samples = class_dist.get_samples(i_class, 400)
    inverted = feature_model.invert(samples)
    mean_bps_per_class.append(
        np.mean(np.abs(np.fft.rfft(var_to_np(inverted.squeeze()))), axis=0))

fig = plt.figure(figsize=(8,3))
plt.plot(np.mean(mean_bps_per_class, axis=1).T)
display_close(fig)

fig = plot_head_signals_tight(np.stack(mean_bps_per_class, axis=-1), sensor_names=sensor_names,
                             figsize=(20,12));
fig = plot_head_signals_tight(np.log(mean_bps_per_class[0]/mean_bps_per_class[1]), sensor_names=sensor_names,
                             figsize=(20,12));
for i_class in range(2):
    samples = class_dist.get_samples(i_class, 20)
    inverted = feature_model.invert(samples)
    plot_head_signals_tight(var_to_np(inverted)[0].squeeze(), sensor_names=sensor_names,
                             figsize=(20,12));
np.array(df.ot_loss_in)[~np.isnan(df.ot_loss_in)]
plt.plot(df.ot_loss_in)
len(df)
df.O<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1267367233232245, 'fraction_numerical': 0.024190704563086255, 'mean_word_length': 2.5889648190331562, 'pattern_counts': {'":': 0, '<': 7, '<?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': 18, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '183588', 'license': 'MIT', 'n_tokens_mistral': 5079, 'n_tokens_neox': 4824, 'n_words': 715, 'path': 'notebooks/bhno-with-adversary/.ipynb_checkpoints/TryAdversary-checkpoint.ipynb', 'repo_name': 'robintibor/reversible2', 'size': '14179'}