<jupyter_start><jupyter_text>Autoencoder Implementation<jupyter_code>class Autoencoder(nn.Module):
  def __init__(self, input_size, latent_size, p_drop = 0.2):
    super(Autoencoder, self).__init__()
    
    self.drop = nn.Dropout(p = p_drop)

    #Define Encoder Layers
    self.fc_e1 = nn.Linear(in_features = input_size, out_features = input_size//8)
    self.fc_e2 = nn.Linear(in_features = input_size//8, out_features = latent_size)

    #Define Decoder Layers
    self.fc_d1 = nn.Linear(in_features = latent_size, out_features = input_size//4)
    self.fc_d2 = nn.Linear(in_features = input_size//4, out_features = input_size)

  def forward(self, input):
    enc = self.drop(nn.Tanh()(self.fc_e1(input)))
    latent = self.fc_e2(enc)
    dec = self.drop(nn.Tanh()(self.fc_d1(latent)))
    output = self.fc_d2(dec)
    return output, latent
X = torch.rand(64, 256)
ae = Autoencoder(input_size= 256, latent_size = 16)
ae(X).shape<jupyter_output><empty_output><jupyter_text>Example: MNIST AutoEncoder<jupyter_code>import torchvision
from math import floor, ceil
batch_size_train = 64 #Define train batch size
batch_size_test  = 256 #Define test batch size (can be larger than train batch size)


# Use the following code to load and normalize the dataset
train_set = torchvision.datasets.MNIST('/files/', train=True, download=True,
                              transform=torchvision.transforms.Compose([
                                torchvision.transforms.ToTensor(),
                                torchvision.transforms.Normalize(
                                  (0.1307,), (0.3081,))
                              ]))


test_set = torchvision.datasets.MNIST('/files/', train=False, download=True,
                             transform=torchvision.transforms.Compose([
                               torchvision.transforms.ToTensor(),
                               torchvision.transforms.Normalize(
                                 (0.1307,), (0.3081,))
                             ]))
train_set[0][0].shape
train_set, val_set = torch.utils.data.random_split(train_set, [floor(len(train_set)*0.8), ceil(len(train_set)*0.2)])
train_dataloader = torch.utils.data.DataLoader(train_set, batch_size = batch_size_train, shuffle = True)
val_dataloader  = torch.utils.data.DataLoader(val_set, batch_size = batch_size_test, shuffle = False)
test_dataloader  = torch.utils.data.DataLoader(test_set, batch_size = batch_size_test, shuffle = False)

max_epochs = 30
learning_rate = 0.001
input_size = 28*28
latent_size = 32
model = Autoencoder(input_size, latent_size)
optimizer = torch.optim.Adam(model.parameters(), lr = learning_rate)
criterion = torch.nn.MSELoss()
from tqdm import tqdm
train_losses = []
val_losses = []
for epoch in tqdm(range(max_epochs)):
  model.train()
  for i, sample in enumerate(train_dataloader):
    optimizer.zero_grad()
    data, label = sample
    data_in = data.view(data.shape[0], input_size)
    output = model(data_in)[0]
    loss = criterion(data_in, output)
    loss.backward()
    train_losses.append(loss.item())
    optimizer.step()
  model.eval()
  vl = 0
  with torch.no_grad():
    for sample in val_dataloader:
      data, label = sample
      data_in = data.view(data.shape[0], input_size)
      output = model(data_in)[0]
      loss = criterion(data_in, output)
      vl += loss.item()
  val_losses.append(vl)
plt.plot(val_losses)
inds = torch.randint(low = 0, size = (8,), high = len(test_set))
samples = 5
f, ax = plt.subplots(samples, 2, figsize = (3, 8))
for i, idx in enumerate(inds[:samples]):
  ax[i,1].imshow(model(test_set[idx][0].flatten().unsqueeze(0))[0].view(28, 28).detach())
  ax[i,0].imshow(test_set[idx][0][0])
f.tight_layout()
test_dataloader = torch.utils.data.DataLoader(test_set, 
                              batch_size = 10000, shuffle = False)
for sample in test_dataloader:
  data = sample[0]
  data_in = data.view(data.shape[0], -1)
  latents = model(data_in)[1]
  labels = sample[1]
U, S, V = torch.pca_lowrank(latents)
plt.figure(figsize = (8,6))
for i in range(10):
  mask = labels == i
  Y = (U[mask]@torch.diag(S)).detach()
  plt.scatter(Y[:,0], Y[:,1], label = i, alpha = 0.4, s = 10)
  plt.legend()<jupyter_output><empty_output><jupyter_text>Assignment: Enhanced Autoencoder<jupyter_code>class EnhancedAutoencoder(nn.Module):
  def __init__(self, input_size, latent_size, p_drop = 0.2):
    super(EnhancedAutoencoder, self).__init__()
    
    self.drop = nn.Dropout(p = p_drop)

    #Define Encoder Layers


    #Define Decoder Layers


  def forward(self, input):
    #Define forward
    return output, latent
# If implementing VAE:
def kld_loss(mus, vars):
  # Define KLD loss with standard normal prior
  return loss<jupyter_output><empty_output><jupyter_text>Sample VAE Implementation (Change this if using for for Assignment)<jupyter_code>class ExampleVAE(nn.Module):
  def __init__(self, input_size, latent_size, p_drop = 0.2):
    super(ExampleVAE, self).__init__()
    
    self.drop = nn.Dropout(p = p_drop)

    #Define Encoder Layers
    self.fc_e1 = nn.Linear(in_features = input_size, out_features = input_size//8)
    self.fc_e2_1 = nn.Linear(in_features = input_size//8, out_features = latent_size)
    self.fc_e2_2 = nn.Linear(in_features = input_size//8, out_features = latent_size)

    #Define Decoder Layers
    self.fc_d1 = nn.Linear(in_features = latent_size, out_features = input_size//4)
    self.fc_d2 = nn.Linear(in_features = input_size//4, out_features = input_size)

  def forward(self, input):
    enc = self.drop(nn.Tanh()(self.fc_e1(input)))
    mus = self.fc_e2_1(enc) #means of latent Gaussian
    var =self.fc_e2_2(enc)  #(log) var of latent Gaussian
    sample = self.reparameterize(mus, var)

    dec = self.drop(nn.Tanh()(self.fc_d1(sample)))
    output = self.fc_d2(dec)
    return output, (mus, var)

  def reparameterize(self, mu, logvar):
    '''
    Used to generate sample from latent Gaussian
    '''
    std = torch.exp(0.5 * logvar)
    eps = torch.randn_like(std)
    return eps * std + mu<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1263418662262593, 'fraction_numerical': 0.02213047068538398, 'mean_word_length': 3.4562178072111847, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '405427', 'license': 'MIT', 'n_tokens_mistral': 2245, 'n_tokens_neox': 2126, 'n_words': 426, 'path': 'labs/10/Lab_10_Autoencoders.ipynb', 'repo_name': 'raph651/amath-563', 'size': '6055'}