6268 enhance hovernet load pretrained function

[yiheng-wang-nv]

Fixes #6268 .

Description

This PR enhances Hovernet's load pretrained function.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[KumoLiu]

LGTM, thanks for update!

[Nic-Ma]

Hi @yiheng-wang-nv @KumoLiu ,

I don't quite understand the use case of this PR, could you please help share some examples here?
I mainly want to ensure the bundle user experience is not affected, especially configs.

Thanks in advance.

[yiheng-wang-nv]

Hi @yiheng-wang-nv @KumoLiu ,

I don't quite understand the use case of this PR, could you please help share some examples here? I mainly want to ensure the bundle user experience is not affected, especially configs.

Thanks in advance.

Hi Nic, some users are used to save not only weights, but also other information such as learning rate, epoch ... into the whole weight file, if print the keys of the file we will see something like:

dict_keys(['epoch', 'arch', 'state_dict', 'optimizer', 'version', 'args', 'amp_scaler', 'metric'])

Where, only the state_dict is the place we need.

Currently, Hovernet is designed to load standard pytorch resnet50 pretrained weights which only contain the weights information, and print the keys of the file we will see:

odict_keys(['conv1.weight', 'bn1.running_mean', 'bn1.running_var', 'bn1.weight', 'bn1.bias', 'layer1.0.conv1.weight', 'layer1.0.bn1.running_mean', 'layer1.0.bn1.running_var', 'layer1.0.bn1.weight', 'layer1.0.bn1.bias', 'layer1.0.conv2.weight', 'layer1.0.bn2.running_mean', 'layer1.0.bn2.running_var', 'layer1.0.bn2.weight', 'layer1.0.bn2.bias', 'layer1.0.conv3.weight', 'layer1.0.bn3.running_mean', 'layer1.0.bn3.running_var', 'layer1.0.bn3.weight', 'layer1.0.bn3.bias', 'layer1.0.downsample.0.weight' ...

[Nic-Ma]

I mean an example to show how to use this arg in a python program and a bundle config.

[yiheng-wang-nv]

I see, in a bundle config we can do:

    "network_def": {
        "_target_": "HoVerNet",
        "mode": "@hovernet_mode",
        "in_channels": 3,
        "out_classes": 5,
        "adapt_standard_resnet": true,
        "pretrained_url": "$None",
        "freeze_encoder": true,
        "pretrained_state_dict_key": "state_dict"
    },

Only add this arg: "pretrained_state_dict_key"

[Nic-Ma]

@yiheng-wang-nv OK, looks good to me, please fix the CI errors.

Thanks.