Best Practice Sub-modules and Hyperparameters

[cheind]

Hey,

what is the best practice when it comes to checkpointing user specified sub-modules passed to __init__ in Lightning? Consider a general network consisting of a backbone net and a variable number of heads:

@dataclass
class HeadSpec:
    name: str
    channels: int
    activation: nn.Module = None

class Model(pl.LightningModule):
    def __init__(
            self,
            backbone: nn.Module = None,
            heads: List[HeadSpec] = None
    ):
        super().__init__()
        self.backbone = backbone
        self.heads = nn.ModuleDict({h.name: make_head(h, chn) for h in heads})
        self.save_hyperparameters()
...

Is the following advised?

net = Model(backbone=ResNet(...), heads=[HeadSpec('scores',1, nn.Sigmoid())])
# ... train and checkpoint net

# later
net = Model.load_from_checkpoint(checkpoint_path="example.ckpt")

Will the above correctly initialize backbone and head structure? This would mean that save_hyperparameters()supports pickling of arbitrary Python objects. According to PyTorch, this is not recommended as it may break your code when at a later point ResNet class changes. Yet, the above strategy seems to be used by LightningFlash.

What's the alternative then? Creating a sub-class for each specific model and making the hyper-parameters simple types? I.e

class MyModel(Model):
    def __init__(self, backbone_type='ResNet'):
        bklass = {
            'ResNet': ResNet,
            'VGG': VGG,
        }[backbone_type]
        super().__init__(backbone=bklass(), heads=[HeadSpec('scores',1, nn.Sigmoid())])
        self.save_hyperparameters()
...

(where I assume now that Model does not save hyperparameters).