# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import os
os.environ['CUDA_VISIBLE_DEVICES'] = "0"
import torch
from torch import nn
from tqdm.auto import tqdm
from matplotlib import pyplot as plt
import seaborn as sns
import numpy as np


class RandomFunction(nn.Module):
    def __init__(self):
        super(RandomFunction, self).__init__()
        self.fc1 = nn.Linear(2, 50)
        self.fc2 = nn.Linear(50, 50)
        self.act = nn.Tanh()
    def forward(self, x):
        x = self.act(self.fc1(x))
        x = self.fc2(x)
        return x


rf = RandomFunction().cuda()

features = torch.randn(1000, 2).cuda()

data = rf(features).detach()


class Decoder(nn.Module):
    def __init__(self):
        super(Decoder, self).__init__()
        self.fc1 = nn.Linear(100, 50)
        self.fc2 = nn.Linear(50, 2)
        self.act = nn.Tanh()
    def forward(self, x):
        x = self.act(self.fc1(x))
        x = self.fc2(x)
        return x



class Switch(nn.Module):
    def __init__(self):
        super(Switch, self).__init__()
        self.probas = nn.Parameter(torch.ones(50, requires_grad=True,
                                                         dtype=torch.float32) * 1.0)
        self.last_mask = None
    def clip_proba(self):
        self.probas.data = torch.clamp(self.probas.data, min=0.5, max=1.0)      
        
    def forward(self, x):
        self.clip_proba()
        
        mask = torch.bernoulli(self.probas.view(1, 50).repeat(x.shape[0], 1)).detach()
        self.last_mask = mask
        return torch.cat([mask, x * mask], dim=1)



sw(data)

dec = Decoder().cuda()
sw = Switch().cuda()

delta = (dec(sw(data)) - features).pow(2).sum(1)


def fill_probas_grad(delta):
    delta_view = delta.view(delta.shape[0], 1).repeat(1, 50)
    loss_on = torch.where(sw.last_mask > 0.5, delta_view, torch.zeros_like(delta_view)).sum(0)
    loss_off = torch.where(sw.last_mask < 0.5, delta_view, torch.zeros_like(delta_view)).sum(0)
    
    cnt_on = sw.last_mask.sum(0)
    cnt_off = (1 - sw.last_mask).sum(0)
    
    cnt_on = torch.clamp(cnt_on, 1, 100000000000000)
    cnt_off = torch.clamp(cnt_off, 1, 1000000000000)
    
    g = (loss_on / cnt_on - loss_off / cnt_off).detach()
    if sw.probas.grad is None:
        sw.probas.grad = g
    else:
        sw.probas.grad.data += g


opt = torch.optim.Adam(list(sw.parameters()) + list(dec.parameters()), lr=1e-3)

losses = []
fits = []
sparsity = []
coeffs = []
coeff = 1e-9

for i in tqdm(range(50000)):
    opt.zero_grad()
    delta = (dec(sw(data)) - features).pow(2).sum(1)
    loss_fit = delta.mean(0)
    
    if i % 50 == 0:
        if loss_fit < 1e-5 and coeff < 10:
            coeff /= 0.5
        elif coeff > 1e-9:
            coeff *= 0.5
    
    loss_sp = sw.probas.abs().sum()
    loss = loss_fit + loss_sp * coeff
    loss.backward()
    fill_probas_grad(delta)
    opt.step()
    
    losses.append(loss.item())
    fits.append(loss_fit.item())
    sparsity.append(loss_sp.item() - 0.5 * 50)
    coeffs.append(coeff)

plt.figure(figsize=(13, 5))
plt.subplot(1, 4, 1)
plt.plot(losses)
plt.title('loss')
plt.yscale('log')
plt.subplot(1, 4, 2)
plt.plot(fits)
plt.title('fit')
plt.yscale('log')
plt.subplot(1, 4, 3)
plt.plot(sparsity)
plt.title('sp')
plt.yscale('log')
plt.subplot(1, 4, 4)
plt.plot(coeffs)
plt.title('coef')
plt.yscale('log')

sns.heatmap(sw.probas.detach().cpu().numpy().reshape(-3, 1),
            vmin=0.5, vmax=1)

sw.probas



StarcoderdataJupyterScriptsDedupFiltered