<jupyter_start><jupyter_text>Dataset<jupyter_code>import sys
sys.path.append('../../datasets/')
from prepare_individuals import prepare, germanBats
import matplotlib.pyplot as plt
import torch
import numpy as np
import tqdm
import pickle

classes = germanBats
patch_len = 44                               # 88 bei 44100, 44 bei 22050 = 250ms ~ 25ms

X_train, Y_train, X_test, Y_test, X_val, Y_val = prepare("../../datasets/prepared.h5", classes, patch_len)
print("Total calls:", len(X_train) + len(X_test) + len(X_val))
print(X_train.shape, Y_train.shape)
X = np.concatenate([X_train, X_test, X_val])
Y = np.concatenate([Y_train, Y_test, Y_val])<jupyter_output><empty_output><jupyter_text>Model<jupyter_code>import time
import datetime
import tqdm
import torch.nn as nn
import torchvision
from torch.utils.data import TensorDataset, DataLoader
class Flatten(nn.Module):
    def forward(self, input):
        return input.view(input.size(0), -1)
class UnFlatten(nn.Module):
    def forward(self, input, size=1024):
        return input.view(input.size(0), size, 1, 1)
class VAE(nn.Module):
    def __init__(self, image_channels=3, h_dim=1024, z_dim=32):
        super(VAE, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(image_channels, 32, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(128, 256, kernel_size=4, stride=2),
            nn.ReLU(),
            Flatten()
        )
        
        self.fc1 = nn.Linear(h_dim, z_dim)
        self.fc2 = nn.Linear(h_dim, z_dim)
        self.fc3 = nn.Linear(z_dim, h_dim)
        
        self.decoder = nn.Sequential(
            UnFlatten(),
            nn.ConvTranspose2d(h_dim, 128, kernel_size=5, stride=2),
            nn.ReLU(),
            nn.ConvTranspose2d(128, 64, kernel_size=5, stride=2),
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, kernel_size=6, stride=2),
            nn.ReLU(),
            nn.ConvTranspose2d(32, image_channels, kernel_size=6, stride=2),
            nn.Sigmoid(),
        )
        
    def reparameterize(self, mu, logvar):
        std = logvar.mul(0.5).exp_()
        # return torch.normal(mu, std)
        esp = torch.randn(*mu.size())
        z = mu + std * esp
        return z
    
    def bottleneck(self, h):
        mu, logvar = self.fc1(h), self.fc2(h)
        z = self.reparameterize(mu, logvar)
        return z, mu, logvar

    def encode(self, x):
        h = self.encoder(x)
        z, mu, logvar = self.bottleneck(h)
        return z, mu, logvar

    def decode(self, z):
        z = self.fc3(z)
        z = self.decoder(z)
        return z

    def forward(self, x):
        z, mu, logvar = self.encode(x)
        z = self.decode(z)
        return z, mu, logvar

def train_epoch(model, epoch, criterion, optimizer, dataloader, device):
    model.train()
    
    running_loss = 0.0
    running_corrects = 0
    
    num_batches = len(dataloader)
    num_samples = len(dataloader.dataset)
    
    for batch, (inputs, _) in enumerate(tqdm.tqdm(dataloader)):
        # Transfer Data to GPU if available
        inputs = inputs.to(device)
        
        # Clear the gradients
        optimizer.zero_grad()
        
        # Forward Pass
        outputs, mu, logvar = model(inputs)
        
        # Compute Loss
        loss, bce, kld = criterion(outputs, inputs, mu, logvar)
            
        # Calculate gradients
        loss.backward()
        
        # Update Weights
        optimizer.step()
        
        running_loss += loss.item() * inputs.size(0)
    
    epoch_loss = running_loss / num_samples
    
    return epoch_loss
from torchsampler import ImbalancedDatasetSampler
from torch.utils.data import WeightedRandomSampler

batch_size = 64
epochs = 50
lr = 0.01
warmup_epochs = 5
wd = 0.01

train_data = TensorDataset(torch.Tensor(np.expand_dims(X, axis=1)), torch.from_numpy(Y))
train_loader = DataLoader(train_data, batch_size=batch_size)
model = VAE(image_channels=1)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    model = nn.DataParallel(model, device_ids=[0, 1])
model.to(device)
print(device)
import wandb

#wandb.init(project="BAT-vae", entity="frankfundel")

wandb.config = {
  "learning_rate": lr,
  "epochs": epochs,
  "batch_size": batch_size
}


def loss_fn(recon_x, x, mu, logvar):
    BCE = F.binary_cross_entropy(recon_x, x, size_average=False)
    KLD = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
    return BCE + KLD


criterion = loss_fn

optimizer = torch.optim.Adam(model.parameters(), lr=lr)

min_val_loss = np.inf

torch.autograd.set_detect_anomaly(True)
for epoch in range(epochs):
    end = time.time()
    print(f"==================== Starting at epoch {epoch} ====================", flush=True)
    
    train_loss, train_acc = train_epoch(model, epoch, criterion, optimizer, train_loader, device)
    print('Training loss: {:.4f} Acc: {:.4f}'.format(train_loss, train_acc), flush=True)
    
    rand_img = train_data[randint(0, len(train_data))].to(device)
    rand_img_pred = model(rand_img)
    images = wandb.Image([rand_img, rand_img_pred], caption="Top: Output, Bottom: Input")
    
    '''wandb.log({
        "train_loss": train_loss,
        "train_acc": train_acc,
        "images": images
    })'''
    
    if min_val_loss > val_loss:
        print('val_loss decreased, saving model', flush=True)
        min_val_loss = val_loss
         
        # Saving State Dict
        torch.save(model.state_dict(), 'vae.pth')
model.load_state_dict(torch.load('vae.pth'))
compiled_model = torch.jit.script(model)
torch.jit.save(compiled_model, 'vae.pt')
wandb.finish()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13378492802709568, 'fraction_numerical': 0.021845893310753598, 'mean_word_length': 2.7593889242520686, 'pattern_counts': {'":': 6, '<': 9, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '283367', 'license': 'MIT', 'n_tokens_mistral': 2171, 'n_tokens_neox': 2077, 'n_words': 435, 'path': 'models/vae/vae.ipynb', 'repo_name': 'FrankFundel/BAT', 'size': '5905'}