With FSDP, when using model's forward in on_fit_end() hook weights are misshappen

[QuentinAndre11]

Hi everyone !
I am training an image classifier and would like to see the embeddings at the end of training, but I don't find how to do it while using FSDP, since the weights seem to get flattenned outside of train/test/validation/_step. Indeed, with the following code, I get a RuntimeError: weight should have at least three dimensions.
I use pytorch v2.1.0 with CUDA 118 and lightning v2.3.3

Thus, how should I do forward calls on_fit_end() hook while using FSDP? Or is there a better way to get the embeddings (and if so, why shouldn't I use the weights outside of training steps)?

Thanks for your help already

import certifi
import os
import timm

import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torch.nn.functional as F
import lightning as L

from lightning.pytorch.core import LightningModule, LightningDataModule
from lightning.pytorch import loggers as pl_loggers
from lightning.pytorch.trainer import Trainer
from lightning.pytorch.strategies import FSDPStrategy

from torchvision.datasets import MNIST
from torchvision.transforms import v2
from torchvision.transforms import Lambda

from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

from einops import rearrange



class ResnetModule(LightningModule):
    def __init__(self, num_classes=10, in_chans=3):
        super().__init__()
        self.model = timm.create_model("resnet50", pretrained = False, num_classes = num_classes, in_chans = in_chans, drop_rate = 0.3)
        self.loss_fn = nn.BCEWithLogitsLoss() 

    def forward(self, x):
        out = self.model(x)
        return out

    def training_step(self, batch, batch_idx):
        # Here we have self.model.conv1.weight.shape = torch.Size([64, 3, 7, 7])
        loss = self._calculate_loss(batch, batch_idx, "train")
        return loss
    
    def validation_step(self, batch, batch_idx):
        # Here we have self.model.conv1.weight.shape = torch.Size([64, 3, 7, 7])
        self._calculate_loss(batch, batch_idx, "val")
        
    def _calculate_loss(self, batch, batch_idx, mode = "train"):
        images, labels = batch
        
        outputs = self(images)
        loss = self.loss_fn(outputs, labels)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.SGD(self.parameters(), 1e-4)
        return optimizer
    
    def get_tb_logger(self, experiment:bool = False) -> pl_loggers.TensorBoardLogger | SummaryWriter:
        for lg in self.trainer.loggers:
            if isinstance(lg, pl_loggers.TensorBoardLogger):
                return lg.experiment if experiment else lg
        return None
    
    def on_train_end(self):
        # !!!!!!!!!
        # Here we have self.model.conv1.weight.shape = torch.Size([9408])
        # !!!!!!!!!
        pass
    
    def on_fit_end(self):
        # !!!!!!!!!
        # Here we have self.model.conv1.weight.shape = torch.Size([9408]) 
        # !!!!!!!!!
        embeddings_activations = []
        embeddings_inputs = []
        embeddings_labels = []
        
        def hook(model, input, output):
            embeddings_activations.append(output)

        # Attach hook to the wanted layer
        hook_handle = self.model.global_pool.register_forward_hook(hook) 
        
        val_dataloader = self.trainer.datamodule.val_dataloader()
        for batch_idx, batch in enumerate(val_dataloader):
            imgs, labels = batch
            embeddings_inputs.append(imgs)
            embeddings_labels.append(labels)
            self(imgs)
        tb_logger = self.get_tb_logger(experiment=True)
        if tb_logger:

            features = rearrange(embeddings_activations, 'n b p -> (n b) p')
            images = rearrange(embeddings_inputs, 'n b c h w -> (n b) c h w')
            labels_one_hot = rearrange(embeddings_labels, 'n b l -> (n b) l')
            metadata = [torch.argmax(t).item() for t in labels_one_hot]
            tb_logger.add_embedding(
                features,
                metadata = metadata,
                label_img = images,
                global_step = self.current_epoch,
                tag = f"{self.model.__class__.__name__}'s embeddings",
            )
            
        hook_handle.remove()
        del embeddings_activations
        del embeddings_inputs
        del embeddings_labels


class MNISTModule(LightningDataModule):
    def __init__(self, num_workers, pin_memory):
        super().__init__()
        self.batch_size = 64
        self.num_workers = num_workers
        self.pin_memory = pin_memory
        self.transform = transforms.Compose(
            [
            v2.Grayscale(num_output_channels=3),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ]
            )
        self.mnist_path = os.path.join(os.getcwd(),'dataset','mnist')
        
    def prepare_data(self):
        dataset = MNIST(
            self.mnist_path,
            download=True,
            transform=self.transform if self.transform is not None else transforms.ToTensor(),
            target_transform=Lambda(lambda y: torch.zeros(10, dtype=torch.float).scatter_(0, torch.tensor(y), value=1))
        )
        
        
    def setup(self, stage: str):
        if stage == "fit":
            self.train_set = MNIST(
                self.mnist_path,
                download=False,
                transform=self.transform if self.transform is not None else transforms.ToTensor(),
                target_transform=Lambda(lambda y: torch.zeros(10, dtype=torch.float).scatter_(0, torch.tensor(y), value=1))
            )
            self.val_set = MNIST(
                self.mnist_path,
                download=False,
                transform=self.transform if self.transform is not None else transforms.ToTensor(),
                target_transform=Lambda(lambda y: torch.zeros(10, dtype=torch.float).scatter_(0, torch.tensor(y), value=1))
            )


    def train_dataloader(self):
        return DataLoader(self.train_set, batch_size=self.batch_size,num_workers=self.num_workers,pin_memory=self.pin_memory)
    
    def val_dataloader(self):
        return DataLoader(self.val_set, batch_size=self.batch_size,num_workers=self.num_workers,pin_memory=self.pin_memory)

def main():
    print("Using PyTorch {} and Lightning {}".format(torch.__version__, L.__version__))

    # Lightning modules    
    datamodule = MNISTModule(num_workers=8, pin_memory=True)
    resnet = ResnetModule(num_classes=10, in_chans=3)
    
    # Logger
    tensorboard = pl_loggers.TensorBoardLogger(
        save_dir = os.path.join(os.getcwd(), 'results'),
        log_graph = False,
        )
    
    # Strategy
    strategy = FSDPStrategy(
            activation_checkpointing_policy={nn.Linear,nn.Conv2d},  
            sharding_strategy="FULL_SHARD",  
        )

    # Trainer
    trainer = Trainer(
        devices=2,
        max_epochs=2,
        strategy=strategy,
        logger=tensorboard,
        )

    trainer.fit(resnet, datamodule)
    trainer.print(torch.cuda.memory_summary())

if __name__ == '__main__':
    main()