|
|
from functools import partial |
|
|
from typing import Any, cast, Dict, List, Optional, Union |
|
|
|
|
|
import torch |
|
|
import torch.nn as nn |
|
|
|
|
|
from ..transforms._presets import ImageClassification |
|
|
from ..utils import _log_api_usage_once |
|
|
from ._api import register_model, Weights, WeightsEnum |
|
|
from ._meta import _IMAGENET_CATEGORIES |
|
|
from ._utils import _ovewrite_named_param, handle_legacy_interface |
|
|
|
|
|
|
|
|
__all__ = [ |
|
|
"VGG", |
|
|
"VGG11_Weights", |
|
|
"VGG11_BN_Weights", |
|
|
"VGG13_Weights", |
|
|
"VGG13_BN_Weights", |
|
|
"VGG16_Weights", |
|
|
"VGG16_BN_Weights", |
|
|
"VGG19_Weights", |
|
|
"VGG19_BN_Weights", |
|
|
"vgg11", |
|
|
"vgg11_bn", |
|
|
"vgg13", |
|
|
"vgg13_bn", |
|
|
"vgg16", |
|
|
"vgg16_bn", |
|
|
"vgg19", |
|
|
"vgg19_bn", |
|
|
] |
|
|
|
|
|
|
|
|
class VGG(nn.Module): |
|
|
def __init__( |
|
|
self, features: nn.Module, num_classes: int = 1000, init_weights: bool = True, dropout: float = 0.5 |
|
|
) -> None: |
|
|
super().__init__() |
|
|
_log_api_usage_once(self) |
|
|
self.features = features |
|
|
self.avgpool = nn.AdaptiveAvgPool2d((7, 7)) |
|
|
self.classifier = nn.Sequential( |
|
|
nn.Linear(512 * 7 * 7, 4096), |
|
|
nn.ReLU(True), |
|
|
nn.Dropout(p=dropout), |
|
|
nn.Linear(4096, 4096), |
|
|
nn.ReLU(True), |
|
|
nn.Dropout(p=dropout), |
|
|
nn.Linear(4096, num_classes), |
|
|
) |
|
|
if init_weights: |
|
|
for m in self.modules(): |
|
|
if isinstance(m, nn.Conv2d): |
|
|
nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") |
|
|
if m.bias is not None: |
|
|
nn.init.constant_(m.bias, 0) |
|
|
elif isinstance(m, nn.BatchNorm2d): |
|
|
nn.init.constant_(m.weight, 1) |
|
|
nn.init.constant_(m.bias, 0) |
|
|
elif isinstance(m, nn.Linear): |
|
|
nn.init.normal_(m.weight, 0, 0.01) |
|
|
nn.init.constant_(m.bias, 0) |
|
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|
|
x = self.features(x) |
|
|
x = self.avgpool(x) |
|
|
x = torch.flatten(x, 1) |
|
|
x = self.classifier(x) |
|
|
return x |
|
|
|
|
|
|
|
|
def make_layers(cfg: List[Union[str, int]], batch_norm: bool = False) -> nn.Sequential: |
|
|
layers: List[nn.Module] = [] |
|
|
in_channels = 3 |
|
|
for v in cfg: |
|
|
if v == "M": |
|
|
layers += [nn.MaxPool2d(kernel_size=2, stride=2)] |
|
|
else: |
|
|
v = cast(int, v) |
|
|
conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1) |
|
|
if batch_norm: |
|
|
layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)] |
|
|
else: |
|
|
layers += [conv2d, nn.ReLU(inplace=True)] |
|
|
in_channels = v |
|
|
return nn.Sequential(*layers) |
|
|
|
|
|
|
|
|
cfgs: Dict[str, List[Union[str, int]]] = { |
|
|
"A": [64, "M", 128, "M", 256, 256, "M", 512, 512, "M", 512, 512, "M"], |
|
|
"B": [64, 64, "M", 128, 128, "M", 256, 256, "M", 512, 512, "M", 512, 512, "M"], |
|
|
"D": [64, 64, "M", 128, 128, "M", 256, 256, 256, "M", 512, 512, 512, "M", 512, 512, 512, "M"], |
|
|
"E": [64, 64, "M", 128, 128, "M", 256, 256, 256, 256, "M", 512, 512, 512, 512, "M", 512, 512, 512, 512, "M"], |
|
|
} |
|
|
|
|
|
|
|
|
def _vgg(cfg: str, batch_norm: bool, weights: Optional[WeightsEnum], progress: bool, **kwargs: Any) -> VGG: |
|
|
if weights is not None: |
|
|
kwargs["init_weights"] = False |
|
|
if weights.meta["categories"] is not None: |
|
|
_ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) |
|
|
model = VGG(make_layers(cfgs[cfg], batch_norm=batch_norm), **kwargs) |
|
|
if weights is not None: |
|
|
model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) |
|
|
return model |
|
|
|
|
|
|
|
|
_COMMON_META = { |
|
|
"min_size": (32, 32), |
|
|
"categories": _IMAGENET_CATEGORIES, |
|
|
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#alexnet-and-vgg", |
|
|
"_docs": """These weights were trained from scratch by using a simplified training recipe.""", |
|
|
} |
|
|
|
|
|
|
|
|
class VGG11_Weights(WeightsEnum): |
|
|
IMAGENET1K_V1 = Weights( |
|
|
url="https://download.pytorch.org/models/vgg11-8a719046.pth", |
|
|
transforms=partial(ImageClassification, crop_size=224), |
|
|
meta={ |
|
|
**_COMMON_META, |
|
|
"num_params": 132863336, |
|
|
"_metrics": { |
|
|
"ImageNet-1K": { |
|
|
"acc@1": 69.020, |
|
|
"acc@5": 88.628, |
|
|
} |
|
|
}, |
|
|
"_ops": 7.609, |
|
|
"_file_size": 506.84, |
|
|
}, |
|
|
) |
|
|
DEFAULT = IMAGENET1K_V1 |
|
|
|
|
|
|
|
|
class VGG11_BN_Weights(WeightsEnum): |
|
|
IMAGENET1K_V1 = Weights( |
|
|
url="https://download.pytorch.org/models/vgg11_bn-6002323d.pth", |
|
|
transforms=partial(ImageClassification, crop_size=224), |
|
|
meta={ |
|
|
**_COMMON_META, |
|
|
"num_params": 132868840, |
|
|
"_metrics": { |
|
|
"ImageNet-1K": { |
|
|
"acc@1": 70.370, |
|
|
"acc@5": 89.810, |
|
|
} |
|
|
}, |
|
|
"_ops": 7.609, |
|
|
"_file_size": 506.881, |
|
|
}, |
|
|
) |
|
|
DEFAULT = IMAGENET1K_V1 |
|
|
|
|
|
|
|
|
class VGG13_Weights(WeightsEnum): |
|
|
IMAGENET1K_V1 = Weights( |
|
|
url="https://download.pytorch.org/models/vgg13-19584684.pth", |
|
|
transforms=partial(ImageClassification, crop_size=224), |
|
|
meta={ |
|
|
**_COMMON_META, |
|
|
"num_params": 133047848, |
|
|
"_metrics": { |
|
|
"ImageNet-1K": { |
|
|
"acc@1": 69.928, |
|
|
"acc@5": 89.246, |
|
|
} |
|
|
}, |
|
|
"_ops": 11.308, |
|
|
"_file_size": 507.545, |
|
|
}, |
|
|
) |
|
|
DEFAULT = IMAGENET1K_V1 |
|
|
|
|
|
|
|
|
class VGG13_BN_Weights(WeightsEnum): |
|
|
IMAGENET1K_V1 = Weights( |
|
|
url="https://download.pytorch.org/models/vgg13_bn-abd245e5.pth", |
|
|
transforms=partial(ImageClassification, crop_size=224), |
|
|
meta={ |
|
|
**_COMMON_META, |
|
|
"num_params": 133053736, |
|
|
"_metrics": { |
|
|
"ImageNet-1K": { |
|
|
"acc@1": 71.586, |
|
|
"acc@5": 90.374, |
|
|
} |
|
|
}, |
|
|
"_ops": 11.308, |
|
|
"_file_size": 507.59, |
|
|
}, |
|
|
) |
|
|
DEFAULT = IMAGENET1K_V1 |
|
|
|
|
|
|
|
|
class VGG16_Weights(WeightsEnum): |
|
|
IMAGENET1K_V1 = Weights( |
|
|
url="https://download.pytorch.org/models/vgg16-397923af.pth", |
|
|
transforms=partial(ImageClassification, crop_size=224), |
|
|
meta={ |
|
|
**_COMMON_META, |
|
|
"num_params": 138357544, |
|
|
"_metrics": { |
|
|
"ImageNet-1K": { |
|
|
"acc@1": 71.592, |
|
|
"acc@5": 90.382, |
|
|
} |
|
|
}, |
|
|
"_ops": 15.47, |
|
|
"_file_size": 527.796, |
|
|
}, |
|
|
) |
|
|
IMAGENET1K_FEATURES = Weights( |
|
|
|
|
|
url="https://download.pytorch.org/models/vgg16_features-amdegroot-88682ab5.pth", |
|
|
transforms=partial( |
|
|
ImageClassification, |
|
|
crop_size=224, |
|
|
mean=(0.48235, 0.45882, 0.40784), |
|
|
std=(1.0 / 255.0, 1.0 / 255.0, 1.0 / 255.0), |
|
|
), |
|
|
meta={ |
|
|
**_COMMON_META, |
|
|
"num_params": 138357544, |
|
|
"categories": None, |
|
|
"recipe": "https://github.com/amdegroot/ssd.pytorch#training-ssd", |
|
|
"_metrics": { |
|
|
"ImageNet-1K": { |
|
|
"acc@1": float("nan"), |
|
|
"acc@5": float("nan"), |
|
|
} |
|
|
}, |
|
|
"_ops": 15.47, |
|
|
"_file_size": 527.802, |
|
|
"_docs": """ |
|
|
These weights can't be used for classification because they are missing values in the `classifier` |
|
|
module. Only the `features` module has valid values and can be used for feature extraction. The weights |
|
|
were trained using the original input standardization method as described in the paper. |
|
|
""", |
|
|
}, |
|
|
) |
|
|
DEFAULT = IMAGENET1K_V1 |
|
|
|
|
|
|
|
|
class VGG16_BN_Weights(WeightsEnum): |
|
|
IMAGENET1K_V1 = Weights( |
|
|
url="https://download.pytorch.org/models/vgg16_bn-6c64b313.pth", |
|
|
transforms=partial(ImageClassification, crop_size=224), |
|
|
meta={ |
|
|
**_COMMON_META, |
|
|
"num_params": 138365992, |
|
|
"_metrics": { |
|
|
"ImageNet-1K": { |
|
|
"acc@1": 73.360, |
|
|
"acc@5": 91.516, |
|
|
} |
|
|
}, |
|
|
"_ops": 15.47, |
|
|
"_file_size": 527.866, |
|
|
}, |
|
|
) |
|
|
DEFAULT = IMAGENET1K_V1 |
|
|
|
|
|
|
|
|
class VGG19_Weights(WeightsEnum): |
|
|
IMAGENET1K_V1 = Weights( |
|
|
url="https://download.pytorch.org/models/vgg19-dcbb9e9d.pth", |
|
|
transforms=partial(ImageClassification, crop_size=224), |
|
|
meta={ |
|
|
**_COMMON_META, |
|
|
"num_params": 143667240, |
|
|
"_metrics": { |
|
|
"ImageNet-1K": { |
|
|
"acc@1": 72.376, |
|
|
"acc@5": 90.876, |
|
|
} |
|
|
}, |
|
|
"_ops": 19.632, |
|
|
"_file_size": 548.051, |
|
|
}, |
|
|
) |
|
|
DEFAULT = IMAGENET1K_V1 |
|
|
|
|
|
|
|
|
class VGG19_BN_Weights(WeightsEnum): |
|
|
IMAGENET1K_V1 = Weights( |
|
|
url="https://download.pytorch.org/models/vgg19_bn-c79401a0.pth", |
|
|
transforms=partial(ImageClassification, crop_size=224), |
|
|
meta={ |
|
|
**_COMMON_META, |
|
|
"num_params": 143678248, |
|
|
"_metrics": { |
|
|
"ImageNet-1K": { |
|
|
"acc@1": 74.218, |
|
|
"acc@5": 91.842, |
|
|
} |
|
|
}, |
|
|
"_ops": 19.632, |
|
|
"_file_size": 548.143, |
|
|
}, |
|
|
) |
|
|
DEFAULT = IMAGENET1K_V1 |
|
|
|
|
|
|
|
|
@register_model() |
|
|
@handle_legacy_interface(weights=("pretrained", VGG11_Weights.IMAGENET1K_V1)) |
|
|
def vgg11(*, weights: Optional[VGG11_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG: |
|
|
"""VGG-11 from `Very Deep Convolutional Networks for Large-Scale Image Recognition <https://arxiv.org/abs/1409.1556>`__. |
|
|
|
|
|
Args: |
|
|
weights (:class:`~torchvision.models.VGG11_Weights`, optional): The |
|
|
pretrained weights to use. See |
|
|
:class:`~torchvision.models.VGG11_Weights` below for |
|
|
more details, and possible values. By default, no pre-trained |
|
|
weights are used. |
|
|
progress (bool, optional): If True, displays a progress bar of the |
|
|
download to stderr. Default is True. |
|
|
**kwargs: parameters passed to the ``torchvision.models.vgg.VGG`` |
|
|
base class. Please refer to the `source code |
|
|
<https://github.com/pytorch/vision/blob/main/torchvision/models/vgg.py>`_ |
|
|
for more details about this class. |
|
|
|
|
|
.. autoclass:: torchvision.models.VGG11_Weights |
|
|
:members: |
|
|
""" |
|
|
weights = VGG11_Weights.verify(weights) |
|
|
|
|
|
return _vgg("A", False, weights, progress, **kwargs) |
|
|
|
|
|
|
|
|
@register_model() |
|
|
@handle_legacy_interface(weights=("pretrained", VGG11_BN_Weights.IMAGENET1K_V1)) |
|
|
def vgg11_bn(*, weights: Optional[VGG11_BN_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG: |
|
|
"""VGG-11-BN from `Very Deep Convolutional Networks for Large-Scale Image Recognition <https://arxiv.org/abs/1409.1556>`__. |
|
|
|
|
|
Args: |
|
|
weights (:class:`~torchvision.models.VGG11_BN_Weights`, optional): The |
|
|
pretrained weights to use. See |
|
|
:class:`~torchvision.models.VGG11_BN_Weights` below for |
|
|
more details, and possible values. By default, no pre-trained |
|
|
weights are used. |
|
|
progress (bool, optional): If True, displays a progress bar of the |
|
|
download to stderr. Default is True. |
|
|
**kwargs: parameters passed to the ``torchvision.models.vgg.VGG`` |
|
|
base class. Please refer to the `source code |
|
|
<https://github.com/pytorch/vision/blob/main/torchvision/models/vgg.py>`_ |
|
|
for more details about this class. |
|
|
|
|
|
.. autoclass:: torchvision.models.VGG11_BN_Weights |
|
|
:members: |
|
|
""" |
|
|
weights = VGG11_BN_Weights.verify(weights) |
|
|
|
|
|
return _vgg("A", True, weights, progress, **kwargs) |
|
|
|
|
|
|
|
|
@register_model() |
|
|
@handle_legacy_interface(weights=("pretrained", VGG13_Weights.IMAGENET1K_V1)) |
|
|
def vgg13(*, weights: Optional[VGG13_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG: |
|
|
"""VGG-13 from `Very Deep Convolutional Networks for Large-Scale Image Recognition <https://arxiv.org/abs/1409.1556>`__. |
|
|
|
|
|
Args: |
|
|
weights (:class:`~torchvision.models.VGG13_Weights`, optional): The |
|
|
pretrained weights to use. See |
|
|
:class:`~torchvision.models.VGG13_Weights` below for |
|
|
more details, and possible values. By default, no pre-trained |
|
|
weights are used. |
|
|
progress (bool, optional): If True, displays a progress bar of the |
|
|
download to stderr. Default is True. |
|
|
**kwargs: parameters passed to the ``torchvision.models.vgg.VGG`` |
|
|
base class. Please refer to the `source code |
|
|
<https://github.com/pytorch/vision/blob/main/torchvision/models/vgg.py>`_ |
|
|
for more details about this class. |
|
|
|
|
|
.. autoclass:: torchvision.models.VGG13_Weights |
|
|
:members: |
|
|
""" |
|
|
weights = VGG13_Weights.verify(weights) |
|
|
|
|
|
return _vgg("B", False, weights, progress, **kwargs) |
|
|
|
|
|
|
|
|
@register_model() |
|
|
@handle_legacy_interface(weights=("pretrained", VGG13_BN_Weights.IMAGENET1K_V1)) |
|
|
def vgg13_bn(*, weights: Optional[VGG13_BN_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG: |
|
|
"""VGG-13-BN from `Very Deep Convolutional Networks for Large-Scale Image Recognition <https://arxiv.org/abs/1409.1556>`__. |
|
|
|
|
|
Args: |
|
|
weights (:class:`~torchvision.models.VGG13_BN_Weights`, optional): The |
|
|
pretrained weights to use. See |
|
|
:class:`~torchvision.models.VGG13_BN_Weights` below for |
|
|
more details, and possible values. By default, no pre-trained |
|
|
weights are used. |
|
|
progress (bool, optional): If True, displays a progress bar of the |
|
|
download to stderr. Default is True. |
|
|
**kwargs: parameters passed to the ``torchvision.models.vgg.VGG`` |
|
|
base class. Please refer to the `source code |
|
|
<https://github.com/pytorch/vision/blob/main/torchvision/models/vgg.py>`_ |
|
|
for more details about this class. |
|
|
|
|
|
.. autoclass:: torchvision.models.VGG13_BN_Weights |
|
|
:members: |
|
|
""" |
|
|
weights = VGG13_BN_Weights.verify(weights) |
|
|
|
|
|
return _vgg("B", True, weights, progress, **kwargs) |
|
|
|
|
|
|
|
|
@register_model() |
|
|
@handle_legacy_interface(weights=("pretrained", VGG16_Weights.IMAGENET1K_V1)) |
|
|
def vgg16(*, weights: Optional[VGG16_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG: |
|
|
"""VGG-16 from `Very Deep Convolutional Networks for Large-Scale Image Recognition <https://arxiv.org/abs/1409.1556>`__. |
|
|
|
|
|
Args: |
|
|
weights (:class:`~torchvision.models.VGG16_Weights`, optional): The |
|
|
pretrained weights to use. See |
|
|
:class:`~torchvision.models.VGG16_Weights` below for |
|
|
more details, and possible values. By default, no pre-trained |
|
|
weights are used. |
|
|
progress (bool, optional): If True, displays a progress bar of the |
|
|
download to stderr. Default is True. |
|
|
**kwargs: parameters passed to the ``torchvision.models.vgg.VGG`` |
|
|
base class. Please refer to the `source code |
|
|
<https://github.com/pytorch/vision/blob/main/torchvision/models/vgg.py>`_ |
|
|
for more details about this class. |
|
|
|
|
|
.. autoclass:: torchvision.models.VGG16_Weights |
|
|
:members: |
|
|
""" |
|
|
weights = VGG16_Weights.verify(weights) |
|
|
|
|
|
return _vgg("D", False, weights, progress, **kwargs) |
|
|
|
|
|
|
|
|
@register_model() |
|
|
@handle_legacy_interface(weights=("pretrained", VGG16_BN_Weights.IMAGENET1K_V1)) |
|
|
def vgg16_bn(*, weights: Optional[VGG16_BN_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG: |
|
|
"""VGG-16-BN from `Very Deep Convolutional Networks for Large-Scale Image Recognition <https://arxiv.org/abs/1409.1556>`__. |
|
|
|
|
|
Args: |
|
|
weights (:class:`~torchvision.models.VGG16_BN_Weights`, optional): The |
|
|
pretrained weights to use. See |
|
|
:class:`~torchvision.models.VGG16_BN_Weights` below for |
|
|
more details, and possible values. By default, no pre-trained |
|
|
weights are used. |
|
|
progress (bool, optional): If True, displays a progress bar of the |
|
|
download to stderr. Default is True. |
|
|
**kwargs: parameters passed to the ``torchvision.models.vgg.VGG`` |
|
|
base class. Please refer to the `source code |
|
|
<https://github.com/pytorch/vision/blob/main/torchvision/models/vgg.py>`_ |
|
|
for more details about this class. |
|
|
|
|
|
.. autoclass:: torchvision.models.VGG16_BN_Weights |
|
|
:members: |
|
|
""" |
|
|
weights = VGG16_BN_Weights.verify(weights) |
|
|
|
|
|
return _vgg("D", True, weights, progress, **kwargs) |
|
|
|
|
|
|
|
|
@register_model() |
|
|
@handle_legacy_interface(weights=("pretrained", VGG19_Weights.IMAGENET1K_V1)) |
|
|
def vgg19(*, weights: Optional[VGG19_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG: |
|
|
"""VGG-19 from `Very Deep Convolutional Networks for Large-Scale Image Recognition <https://arxiv.org/abs/1409.1556>`__. |
|
|
|
|
|
Args: |
|
|
weights (:class:`~torchvision.models.VGG19_Weights`, optional): The |
|
|
pretrained weights to use. See |
|
|
:class:`~torchvision.models.VGG19_Weights` below for |
|
|
more details, and possible values. By default, no pre-trained |
|
|
weights are used. |
|
|
progress (bool, optional): If True, displays a progress bar of the |
|
|
download to stderr. Default is True. |
|
|
**kwargs: parameters passed to the ``torchvision.models.vgg.VGG`` |
|
|
base class. Please refer to the `source code |
|
|
<https://github.com/pytorch/vision/blob/main/torchvision/models/vgg.py>`_ |
|
|
for more details about this class. |
|
|
|
|
|
.. autoclass:: torchvision.models.VGG19_Weights |
|
|
:members: |
|
|
""" |
|
|
weights = VGG19_Weights.verify(weights) |
|
|
|
|
|
return _vgg("E", False, weights, progress, **kwargs) |
|
|
|
|
|
|
|
|
@register_model() |
|
|
@handle_legacy_interface(weights=("pretrained", VGG19_BN_Weights.IMAGENET1K_V1)) |
|
|
def vgg19_bn(*, weights: Optional[VGG19_BN_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG: |
|
|
"""VGG-19_BN from `Very Deep Convolutional Networks for Large-Scale Image Recognition <https://arxiv.org/abs/1409.1556>`__. |
|
|
|
|
|
Args: |
|
|
weights (:class:`~torchvision.models.VGG19_BN_Weights`, optional): The |
|
|
pretrained weights to use. See |
|
|
:class:`~torchvision.models.VGG19_BN_Weights` below for |
|
|
more details, and possible values. By default, no pre-trained |
|
|
weights are used. |
|
|
progress (bool, optional): If True, displays a progress bar of the |
|
|
download to stderr. Default is True. |
|
|
**kwargs: parameters passed to the ``torchvision.models.vgg.VGG`` |
|
|
base class. Please refer to the `source code |
|
|
<https://github.com/pytorch/vision/blob/main/torchvision/models/vgg.py>`_ |
|
|
for more details about this class. |
|
|
|
|
|
.. autoclass:: torchvision.models.VGG19_BN_Weights |
|
|
:members: |
|
|
""" |
|
|
weights = VGG19_BN_Weights.verify(weights) |
|
|
|
|
|
return _vgg("E", True, weights, progress, **kwargs) |
|
|
|
