Skip to content

Enhance Hovernet load pretrained weights functionality #6268

@yiheng-wang-nv

Description

@yiheng-wang-nv

Currently, _remap_standard_resnet_model is designed according to the network structure shown in https://pytorch.org/vision/main/models/generated/torchvision.models.resnet50.html .
Therefore, we assume that there is not extra key (see also here ) to extract the expected state dict.
It would be nice to make the network be compatible with the weights that contain some other keys. For example, users may also save other information rather than state dict within the weights:

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

Therefore, I suggest we can add an arg to support this kind of case. In addition, we can also add some notifications within the load function to help users double check that if the load method is successful.

cc @KumoLiu

Metadata

Metadata

Labels

No labels
No labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions