Custom weight initialization using UNET

[jarod1212]

Hi, I'm trying to initialize my model with a simple trained autoencoder in order that my model can already be familiar with my type of data. I then want to train it with a MONAI BasicUnet to perform binary segmentation.

The segmentation part is working but is there a way to implement the custom initialization keeping a MONAI core structure in my code ?

Here is my basic init function in my network class :

    def __init__(self):
        super().__init__()
        
        # Basic Unet for our model
        self._model = BasicUNet(
            spatial_dims=2,
            in_channels=1,
            out_channels=2, 
        )

Thanks a lot !!

[KumoLiu]

Hi @jarod1212, if your autoencoder has the same structure as the BasicUnet, you can simply load the pre-trained weight by adding a load_weight func like:

def load_weight(self._model, weights_from_autoencoder):
    self._model.load_state_dict(weights_from_autoencoder)

Hope it can help you, thanks!

[jarod1212]

Thank you for your swift answer @KumoLiu ! It helps a lot indeed.