baseline.py

import torch
import torch.nn as nn


class AE(nn.Module):
    def __init__(self,gene_size, latent_dim):
        super(AE, self).__init__()

        # Encoder
        self.encoder = nn.Sequential(
            nn.Linear(gene_size, 1024),
            nn.ReLU(True),
            nn.Linear(1024, 512),
            nn.ReLU(True),
            nn.Linear(512, latent_dim),
            nn.ReLU(True)
        )

        # Decoder
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 512),
            nn.ReLU(True),
            nn.Linear(512, 1024),
            nn.ReLU(True),
            nn.Linear(1024, gene_size),
            nn.Sigmoid()  # or nn.Tanh(), depending on data preprocessing
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x
    
    def latent(self, x):
        x = self.encoder(x)
        return x