-
Notifications
You must be signed in to change notification settings - Fork 298
/
api.py
107 lines (80 loc) · 2.86 KB
/
api.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
"""APIs for querying and loading pre-trained model weights.
See the following references for design details:
* https://pytorch.org/blog/easily-list-and-initialize-models-with-new-apis-in-torchvision/
* https://pytorch.org/vision/stable/models.html
* https://github.com/pytorch/vision/blob/main/torchvision/models/_api.py
""" # noqa: E501
from collections.abc import Callable
from typing import Any
import torch.nn as nn
from torchvision.models._api import WeightsEnum
from .dofa import (
DOFABase16_Weights,
DOFALarge16_Weights,
dofa_base_patch16_224,
dofa_large_patch16_224,
)
from .resnet import ResNet18_Weights, ResNet50_Weights, resnet18, resnet50
from .swin import Swin_V2_B_Weights, swin_v2_b
from .vit import ViTSmall16_Weights, vit_small_patch16_224
_model = {
'dofa_base_patch16_224': dofa_base_patch16_224,
'dofa_large_patch16_224': dofa_large_patch16_224,
'resnet18': resnet18,
'resnet50': resnet50,
'swin_v2_b': swin_v2_b,
'vit_small_patch16_224': vit_small_patch16_224,
}
_model_weights = {
dofa_base_patch16_224: DOFABase16_Weights,
dofa_large_patch16_224: DOFALarge16_Weights,
resnet18: ResNet18_Weights,
resnet50: ResNet50_Weights,
swin_v2_b: Swin_V2_B_Weights,
vit_small_patch16_224: ViTSmall16_Weights,
'dofa_base_patch16_224': DOFABase16_Weights,
'dofa_large_patch16_224': DOFALarge16_Weights,
'resnet18': ResNet18_Weights,
'resnet50': ResNet50_Weights,
'swin_v2_b': Swin_V2_B_Weights,
'vit_small_patch16_224': ViTSmall16_Weights,
}
def get_model(name: str, *args: Any, **kwargs: Any) -> nn.Module:
"""Get an instantiated model from its name.
.. versionadded:: 0.4
Args:
name: Name of the model.
*args: Additional arguments passed to the model builder method.
**kwargs: Additional keyword arguments passed to the model builder method.
Returns:
An instantiated model.
"""
model: nn.Module = _model[name](*args, **kwargs)
return model
def get_model_weights(name: Callable[..., nn.Module] | str) -> WeightsEnum:
"""Get the weights enum class associated with a given model.
.. versionadded:: 0.4
Args:
name: Model builder function or the name under which it is registered.
Returns:
The weights enum class associated with the model.
"""
return _model_weights[name]
def get_weight(name: str) -> WeightsEnum:
"""Get the weights enum value by its full name.
.. versionadded:: 0.4
Args:
name: Name of the weight enum entry.
Returns:
The requested weight enum.
"""
return eval(name)
def list_models() -> list[str]:
"""List the registered models.
.. versionadded:: 0.4
Returns:
A list of registered models.
"""
return list(_model.keys())