<jupyter_start><jupyter_text>Data<jupyter_code>DATA_PATH = 'data/mnist.pkl'

if not os.path.exists(DATA_PATH):
    X, y = fetch_openml('mnist_784', version=1, return_X_y=True, data_home='./data')
    X, y = X.astype(np.float32).reshape(-1, 28, 28), y.astype(np.int64).reshape(-1, 1)
    X = X / X.max()
    
    with open(DATA_PATH, 'wb') as f:
        pickle.dump((X, y), f)
else:
    with open(DATA_PATH, 'rb') as f:
        X, y = pickle.load(f)
X = X.reshape(-1, 784)
class Encoder(nn.Module):

    def __init__(self, latent_dim=10, dropout=.1):
        
        super(Encoder, self).__init__()
        
        self.net = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(784, 128),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(128, 128),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(128, 128),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(64, latent_dim),
        )
        
    def forward(self, x):
        return self.net(x)
class Decoder(nn.Module):

    def __init__(self, latent_dim=10, dropout=.1):
        
        super(Decoder, self).__init__()
        
        self.net = nn.Sequential(
            nn.Linear(latent_dim, 64),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(64, 128),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(128, 128),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(128, 128),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(128, 784),
        )
        
    def forward(self, x):
        return torch.sigmoid(self.net(x))
class AutoEncoder(nn.Module):

    def __init__(self, latent_dim=10, dropout=.1):
        
        super(AutoEncoder, self).__init__()
        
        self.encoder = Encoder(latent_dim, dropout)
        self.decoder = Decoder(latent_dim, dropout)
        
    def forward(self, x):
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        
        return decoded, encoded
plt.imshow(X[0].reshape(28, 28))<jupyter_output><empty_output><jupyter_text>Vanilla<jupyter_code>class AutoEncoderNet(NeuralNetRegressor):
    
    def encode(self, x):
        return self.module_.encoder(torch.tensor(x))
    
    
    def decode(self, x):
        return self.module_.decoder(torch.tensor(x))
vanilla = AutoEncoderNet(
    module=AutoEncoder,
    module__latent_dim=2,
    optimizer=torch.optim.Adam,
    criterion=AutoEncoderLoss,
    criterion__reconstruction_criterion=nn.MSELoss,
    max_epochs=10,
    lr=-1e-3,
    verbose=0,
    callbacks=[
        ('checkpoint', Checkpoint()),
        ('early_stopping', EarlyStopping(patience=3)),
    ]
)
params = dict(
    lr=loguniform(1e-5, 1e-2),
    module__dropout=loguniform(1e-2, .5),
)

cv = RandomizedSearchCV(vanilla, param_distributions=params, n_iter=5, cv=3, verbose=2, n_jobs=1)
vanilla.fit(X, X)
vanilla.decode(vanilla.encode(X)).shape
class L1Penalty(nn.Module):
    
    def __init__(self, ksi=1.):
        super(L1Penalty, self).__init__()
        self.ksi = ksi
        
    def forward(self, code):
        return torch.abs(code).mean()
sparse = NeuralNetRegressor(
    module=VariationalAutoEncoder,
    module__latent_dim=2,
    optimizer=torch.optim.Adam,
    criterion=AutoEncoderLoss,
    criterion__reconstruction_criterion=nn.MSELoss,
    criterion__code_criterion=L1Penalty,
    criterion__code_kwargs=dict(ksi=1e-3),
    max_epochs=30,
    lr=1e-4,
    callbacks=[
        Checkpoint(),
        EarlyStopping(patience=3),
    ]
)
class VariationalAutoEncoder(nn.Module):

    def __init__(self, latent_dim=10, dropout=.1):
        
        super(VariationalAutoEncoder, self).__init__()
        
        self.encoder = Encoder(latent_dim * 2, dropout)
        self.decoder = Decoder(latent_dim, dropout)
        
    def sample(self, mu, log_sigma):
        """
        Samples a vector within the latent space.

        Args:
            mu (torch.Tensor): The mean of the sampled tensor.
            log_sigma (torch.Tensor): The log-variance of the sampled tensor (log for stability reasons).

        Returns:
            z (torch.Tensor): The latent stochastic encoding.
        """

        sigma = torch.exp(.5 * log_sigma)
        epsilon = torch.randn_like(sigma)

        z = mu + sigma * epsilon

        return z
        
    def forward(self, x):
        encoded = self.encoder(x)
        
        code = torch.chunk(encoded, 2, dim=-1)
        
        sample = self.sample(*code)
        decoded = self.decoder(sample)
        
        return code, decoded
class KullbackLeiblerPenalty(nn.Module):
    
    def __init__(self, ksi=1.):
        super(KullbackLeiblerPenalty, self).__init__()
        self.ksi = ksi
        
    def forward(self, code):
        mu, log_sigma = code
        kl = -0.5 * (1 + log_sigma - mu.pow(2) - log_sigma.exp()).sum()
        return self.ksi * kl
vae = NeuralNetRegressor(
    module=VariationalAutoEncoder,
    module__latent_dim=2,
    optimizer=torch.optim.RMSprop,
    criterion=AutoEncoderLoss,
    criterion__reconstruction_criterion=nn.MSELoss,
    criterion__code_criterion=KullbackLeiblerPenalty,
    criterion__code_kwargs=dict(ksi=.1),
    max_epochs=30,
    lr=1e-3,
    callbacks=[
        Checkpoint(),
        EarlyStopping(patience=3),
    ]
)
vae.fit(X, X)
vanilla.fit(X, X)
grid = np.stack([
    np.linspace(-1, 1, 21).reshape(-1, 1).repeat(21, axis=1),
    np.linspace(-1, 1, 21).reshape(1, -1).repeat(21, axis=0),
], axis=2)
grid = torch.tensor(grid.reshape(-1, 2), dtype=torch.float) * 1
reconstructed_grid = cv.module_.decoder(grid).detach().numpy()
reconstructed_grid = reconstructed_grid.reshape(21, 21, 28, 28)
reconstructed_grid = np.concatenate([
    np.concatenate([rg__ for rg__ in rg_], axis=1)
    for rg_ in reconstructed_grid
], axis=0)
plt.figure(figsize=(10, 10))
plt.imshow(reconstructed_grid)
grid.shape
grid = grid.reshape(-1, 2)
X
code, reconstruction = vae.forward(X)
reconstruction = torch.sigmoid(reconstruction)
i = 0
i += 1
plt.imshow(
    np.concatenate(
        [X[i].reshape(28, 28), reconstruction[i].reshape(28, 28)],
        axis=1,
    )
)
from sklearn.manifold import TSNE
tsne = TSNE()
subsample = np.random.choice(np.arange(len(y)), size=1000, replace=False)
transformed = tsne.fit_transform(code[subsample])
plt.figure(figsize=(10, 8))
for n in range(10):
    indices = y[subsample].reshape(-1) == n
    plt.scatter(*transformed[indices].T, color=f'C{n}', alpha=.5)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1348802395209581, 'fraction_numerical': 0.02874251497005988, 'mean_word_length': 2.525593667546174, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, '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': '227793', 'license': 'MIT', 'n_tokens_mistral': 2519, 'n_tokens_neox': 2326, 'n_words': 413, 'path': 'notebooks/Vanilla Autoencoder.ipynb', 'repo_name': 'bdura/bestiary', 'size': '6680'}