Skip to content

Extend TorchVision's existing model builders to support multiple weights #4611

Closed
@datumbox

Description

@datumbox

🚀 The feature

Note: To track the progress of the project check out this board.

The existing model builders support only single pre-trained weights.

To specify if the pretrained weights should be loaded, we pass a boolean pretrained parameter. Example:

def resnet50(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:

Builders that require model composition (for example in Object Detection), receive an additional extra boolean parameter called pretrained_backbone which indicates if a pre-trained backbone should be loaded:

def ssd300_vgg16(
pretrained: bool = False,
progress: bool = True,
num_classes: int = 91,
pretrained_backbone: bool = True,
trainable_backbone_layers: Optional[int] = None,
**kwargs: Any,
):

The above approach is limiting and doesn't allow us to support multiple weights for the same model builder.

Motivation, pitch

The Batteries Included project, enabled us to train far more accurate models than before. In order to make the new weights available to the users, we should extend the current model builders to support multiple weights. This can be done by converting the boolean parameters to an object that contains information about the weights.

Alternatives

An alternative approach would be to modify the weight URLs in place (replace the old with new ones). Unfortunately this would break the BC guarantees of vision.

Additional context

There was an original RFC with far bigger scope that discussed Model Versioning and included a solution for multi-pretrained weights support. We've decided to separate this vision specific extension from the original RFC to cover our needs and give additional time until the model versioning problem is thoroughly investigated.

Metadata

Metadata

Assignees

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions