pytorch-lightning output embeddings completely differnt than pytorch vanilla

[AMA111]

Hey pytorch-lightning team,

I am trying to build an autoencoder architecture via pytorch-lightning (newbie), which I observe its advantages while running complex models in comparison to pytorch implementation.

However in this AE vanilla model, I observe significant difference in the embedding output of pytorch-lightning vs pytorch where I expect that pytorch-lightning should give the exact UMAP embedding as in pytorch

and here is an example:

Pytorch implementation

## AE Architecture
class Autoencoder(nn.Module):
    def __init__(self):
        super().__init__()        
        self.encoder = nn.Sequential(
            #nn.Flatten(),
            nn.Linear(nfeatures_rna, 200),
            nn.LeakyReLU(negative_slope=0.01),
            nn.BatchNorm1d(200, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True),
            nn.Linear(200, 10), # (N, 784) -> (N, 128)
            nn.Dropout(0.2)
        )
        
        self.decoder = nn.Sequential(
            nn.Linear(10, 200), 
            nn.LeakyReLU(negative_slope=0.01),
            nn.BatchNorm1d(200, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True),
            nn.Linear(200, nfeatures_rna), # (N, 784) -> (N, 128)
        )
    def forward(self, x):
        x = torch.flatten(x, start_dim=1)
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        return decoded
modelV = Autoencoder()
lr = 1e-2
optimizer = torch.optim.Adam(modelV.parameters(), lr=lr)
num_epochs = 50
outputs = []
for epoch in range(num_epochs):
    for x, y in data_loader:
        #x = x.reshape(-1, nfeatures_rna) # -> use for Autoencoder_Linear
        recon = modelV(x)
        
        loss = criterion(recon, x)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print(f'Epoch:{epoch+0}, Loss:{loss.item():.4f}')
    outputs.append((epoch, x, recon))
def get_encodings(model, dl):
    #device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.eval()
    with torch.no_grad():
        encodings = [model.encoder(x) for x, _ in dl]
    return torch.cat(encodings, dim=0)
encodings = get_encodings(modelV, data_loader)
encodings = encodings.cpu().numpy()
encodings.shape
## UMAP
import umap
embedding = umap.UMAP(random_state=0).fit_transform(encodings)
plot_df = pd.DataFrame({'A' : []})
#plot_df = metadata.copy()
plot_df["UMAP1"] = embedding[:, 0]
plot_df["UMAP2"] = embedding[:, 1]

Here is the UMAP plot for pytorch embedding:

Pytorch-lightning implementation for the same Pytorch AE model above

import pytorch_lightning as pl
import torch
from torch import nn
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
lr = 1e-2
class Autoencoder(pl.LightningModule):
    def __init__(self):
        super().__init__()        
        
        ## Encoder Arch
        
        self.encoder = nn.Sequential(
            # nn.Flatten(),
            nn.Linear(nfeatures_rna, 200),
            nn.LeakyReLU(negative_slope=0.01),
            nn.BatchNorm1d(200, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True),
            nn.Linear(200, 10), # (N, 784) -> (N, 128)
            nn.Dropout(0.2)
        )
        
        ## Decoder Arch
        
        self.decoder = nn.Sequential(
            nn.Linear(10, 200), 
            nn.LeakyReLU(negative_slope=0.01),
            nn.BatchNorm1d(200, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True),
            nn.Linear(200, nfeatures_rna), # (N, 784) -> (N, 128)
        )
    
    ## Traning step
    
    def training_step(self, batch, batch_idx):
        x, y = batch
        # encode
        #x = x.view(x.size(0), -1)
        z = self.encoder(x)
    
    
        # decode
        recons = self.decoder(z)
        
        # reconstruction #1#
        reconstruction_loss_1 = nn.functional.mse_loss(x, recons)
        return reconstruction_loss_1
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=lr)
    
## Fotward step
    
    def forward(self, x):
        z = self.encoder(x)
        recons = self.decoder(z)
        embedding = z
        return embedding 
autoencoder = Autoencoder()
trainer = pl.Trainer(gpus=1, max_epochs=50)
# trainer = pl.Trainer(gpus=1)
trainer.fit(autoencoder, data_loader)
def get_encodings(model, dl):
    #device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.eval()
    with torch.no_grad():
        encodings = [model.encoder(x) for x, _ in dl]
    return torch.cat(encodings, dim=0)
model = Autoencoder()
encodings = get_encodings(model, data_loader)
encodings = encodings.cpu().numpy()
encodings.shape
embedding = umap.UMAP(random_state=0).fit_transform(encodings) 
plot_df = pd.DataFrame({'A' : []})
plot_df["UMAP1"] = embedding[:, 0]
plot_df["UMAP2"] = embedding[:, 1]
##
custom = sns.scatterplot(data=plot_df, x='UMAP1', 
                      y='UMAP2', 
                      #hue='Gen_and_RealCells',
                      palette="deep", 
                      #style="Gen_and_RealCells"
                     )
plt.legend(bbox_to_anchor=(1.02, 1), loc='upper left', borderaxespad=0)

Here is the UMAP plot for pytorch-lightning embedding :

Since I have experience with this data type and the data generative process, which give me confidence to believe in the pytorch output

Any hints or explanations would be highly appreciated! Thank you in advance for your help.

Best wishes,
Abdelrahman

[rohitgr7]

Hi
here, I think you are reinitializing the model and using the one with random weights.

model = Autoencoder()
encodings = get_encodings(model, data_loader)

you need to load a checkpoint here

model = AutoEncoder.load_from_checkpoint(ckpt_path)
encodings = get_encodings(model, data_loader)

[AMA111]

Hey,

Yes, this worked:

trainer.save_checkpoint("model_0.ckpt")
model = autoencoder.load_from_checkpoint("model_0.ckpt")
encodings = get_encodings(model, data_loader)
encodings = encodings.cpu().numpy()
encodings.shape

Thank you :)

Abdelrahman