VAE reconstruction

Hi all. I am running a VAE on CNN. It does correctly perform reconstruction. But I do fail in flagging the anomalies since the reconstruction lossi is performed as:

def loss_function(self, x, recon, mu, logvar):

          recon_loss = F.smooth_l1_loss(recon, x, reduction='mean') #mean = anomalies diluted across features, sum = amplifies spikes

          kld = -0.5 * torch.mean(1 + logvar - mu ** 2 - logvar.exp()) #KL divergence

          return recon_loss, kld  

This of course sums over all the embedding dimensions (4) where the only one i need is the one of value:

cat_features = cat_vars



embed_cats = [7,2] #get_cat_cardinalities(ds) #[7,2] #get unique



hparams_cnn = OrderedDict(

    run='cnn_emb2',

    cont_vars = cont_vars,

    cat_vars = cat_vars,

    embedding_sizes = [(embed_cats[i], 4) for i in range(len(embed_cats))], #(7,16), (2,16)

    latent_dim = 4,



    #CNN

    kernel = 3,

    layer_sizes = '64,128', #conv layers

    padding = 1,

    batch_norm = True,

    window = 1,



    #Optimizer & Train

    stdev = 0.1,

    kld_beta = 0.05,

    lr = 0.001,

    weight_decay = 1e-5,

    batch_size = 128,

    epochs = 30,

    train_supply_ids = ["SUPPLY011","SUPPLY012","SUPPLY013"]

) 

what to do ?

This is how I reconstruct it:

#inference

trained_model = VAE.load_from_checkpoint('vae_weights.ckpt')

trained_model.freeze()




#move inputs to same device:

device = next(trained_model.parameters()).device




losses = []

all_x = []

all_recon = []




#inference

with torch.no_grad():

    for i in range(len(ds_test_reg)):

        x_cont, x_cat = ds_test_reg[i] #in dataset order




        x_cont = x_cont.unsqueeze(0).unsqueeze(0).to(device)

        x_cat = x_cat.unsqueeze(0).unsqueeze(0).to(device)




        recon, mu, logvar, x = trained_model.forward((x_cont, x_cat))




        #append results

        all_x.append(x.cpu())

        all_recon.append(recon.cpu())




        #loss = recon + KLD

        recon_loss, kld = trained_model.loss_function(x, recon, mu, logvar)

        loss = recon_loss + trained_model.hparams.kld_beta * kld

        losses.append(loss.item()) #moving it to CPU




#x & reconstruction

all_x = torch.cat(all_x, dim=0) #structure [32 | 2 dims (t and val)]

all_recon = torch.cat(all_recon, dim=0)




#Add loss to test_reg

data_with_losses = ds_test_reg.df

data_with_losses['loss'] = np.asarray(losses) #add losses column

data_with_losses.tail()