Skip to content

GitLab

Unverified Commit 5f0edb97 authored by Philip Meier's avatar Philip Meier Committed by GitHub
Browse files

Add ufmt (usort + black) as code formatter (#4384) 



* add ufmt as code formatter

* cleanup

* quote ufmt requirement

* split imports into more groups

* regenerate circleci config

* fix CI

* clarify local testing utils section

* use ufmt pre-commit hook

* split relative imports into local category

* Revert "split relative imports into local category"

This reverts commit f2e224cde2008c56c9347c1f69746d39065cdd51.

* pin black and usort dependencies

* fix local test utils detection

* fix ufmt rev

* add reference utils to local category

* fix usort config

* remove custom categories sorting

* Run pre-commit without fixing flake8

* got a double import in merge
Co-authored-by: default avatarNicolas Hug <nicolashug@fb.com>
parent e45489b1
Hide whitespace changes
Inline Side-by-side
Showing with 569 additions and 525 deletions
torchvision/models/quantization/shufflenetv2.py
View file @ 5f0edb97
from typing import Any
import torch
import torch.nn as nn
from torch import Tensor
from typing import Any
from torchvision.models import shufflenetv2
from ..._internally_replaced_utils import load_state_dict_from_url
from torchvision.models import shufflenetv2
from .utils import _replace_relu, quantize_model
__all__ = [
'QuantizableShuffleNetV2', 'shufflenet_v2_x0_5', 'shufflenet_v2_x1_0',
'shufflenet_v2_x1_5', 'shufflenet_v2_x2_0'
"QuantizableShuffleNetV2",
"shufflenet_v2_x0_5",
"shufflenet_v2_x1_0",
"shufflenet_v2_x1_5",
"shufflenet_v2_x2_0",
]
quant_model_urls = {
'shufflenetv2_x0.5_fbgemm': None,
'shufflenetv2_x1.0_fbgemm':
'https://download.pytorch.org/models/quantized/shufflenetv2_x1_fbgemm-db332c57.pth',
'shufflenetv2_x1.5_fbgemm': None,
'shufflenetv2_x2.0_fbgemm': None,
"shufflenetv2_x0.5_fbgemm": None,
"shufflenetv2_x1.0_fbgemm": "https://download.pytorch.org/models/quantized/shufflenetv2_x1_fbgemm-db332c57.pth",
"shufflenetv2_x1.5_fbgemm": None,
"shufflenetv2_x2.0_fbgemm": None,
}
......@@ -42,9 +45,7 @@ class QuantizableShuffleNetV2(shufflenetv2.ShuffleNetV2):
# TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659
def __init__(self, *args: Any, **kwargs: Any) -> None:
super(QuantizableShuffleNetV2, self).__init__( # type: ignore[misc]
*args,
inverted_residual=QuantizableInvertedResidual,
**kwargs
*args, inverted_residual=QuantizableInvertedResidual, **kwargs
)
self.quant = torch.quantization.QuantStub()
self.dequant = torch.quantization.DeQuantStub()
......@@ -69,9 +70,7 @@ class QuantizableShuffleNetV2(shufflenetv2.ShuffleNetV2):
for m in self.modules():
if type(m) == QuantizableInvertedResidual:
if len(m.branch1._modules.items()) > 0:
torch.quantization.fuse_modules(
m.branch1, [["0", "1"], ["2", "3", "4"]], inplace=True
)
torch.quantization.fuse_modules(m.branch1, [["0", "1"], ["2", "3", "4"]], inplace=True)
torch.quantization.fuse_modules(
m.branch2,
[["0", "1", "2"], ["3", "4"], ["5", "6", "7"]],
......@@ -93,19 +92,18 @@ def _shufflenetv2(
if quantize:
# TODO use pretrained as a string to specify the backend
backend = 'fbgemm'
backend = "fbgemm"
quantize_model(model, backend)
else:
assert pretrained in [True, False]
if pretrained:
if quantize:
model_url = quant_model_urls[arch + '_' + backend]
model_url = quant_model_urls[arch + "_" + backend]
else:
model_url = shufflenetv2.model_urls[arch]
state_dict = load_state_dict_from_url(model_url,
progress=progress)
state_dict = load_state_dict_from_url(model_url, progress=progress)
model.load_state_dict(state_dict)
return model
......@@ -127,8 +125,9 @@ def shufflenet_v2_x0_5(
progress (bool): If True, displays a progress bar of the download to stderr
quantize (bool): If True, return a quantized version of the model
"""
return _shufflenetv2('shufflenetv2_x0.5', pretrained, progress, quantize,
[4, 8, 4], [24, 48, 96, 192, 1024], **kwargs)
return _shufflenetv2(
"shufflenetv2_x0.5", pretrained, progress, quantize, [4, 8, 4], [24, 48, 96, 192, 1024], **kwargs
)
def shufflenet_v2_x1_0(
......@@ -147,8 +146,9 @@ def shufflenet_v2_x1_0(
progress (bool): If True, displays a progress bar of the download to stderr
quantize (bool): If True, return a quantized version of the model
"""
return _shufflenetv2('shufflenetv2_x1.0', pretrained, progress, quantize,
[4, 8, 4], [24, 116, 232, 464, 1024], **kwargs)
return _shufflenetv2(
"shufflenetv2_x1.0", pretrained, progress, quantize, [4, 8, 4], [24, 116, 232, 464, 1024], **kwargs
)
def shufflenet_v2_x1_5(
......@@ -167,8 +167,9 @@ def shufflenet_v2_x1_5(
progress (bool): If True, displays a progress bar of the download to stderr
quantize (bool): If True, return a quantized version of the model
"""
return _shufflenetv2('shufflenetv2_x1.5', pretrained, progress, quantize,
[4, 8, 4], [24, 176, 352, 704, 1024], **kwargs)
return _shufflenetv2(
"shufflenetv2_x1.5", pretrained, progress, quantize, [4, 8, 4], [24, 176, 352, 704, 1024], **kwargs
)
def shufflenet_v2_x2_0(
......@@ -187,5 +188,6 @@ def shufflenet_v2_x2_0(
progress (bool): If True, displays a progress bar of the download to stderr
quantize (bool): If True, return a quantized version of the model
"""
return _shufflenetv2('shufflenetv2_x2.0', pretrained, progress, quantize,
[4, 8, 4], [24, 244, 488, 976, 2048], **kwargs)
return _shufflenetv2(
"shufflenetv2_x2.0", pretrained, progress, quantize, [4, 8, 4], [24, 244, 488, 976, 2048], **kwargs
)
torchvision/models/quantization/utils.py
View file @ 5f0edb97
......@@ -23,14 +23,15 @@ def quantize_model(model: nn.Module, backend: str) -> None:
torch.backends.quantized.engine = backend
model.eval()
# Make sure that weight qconfig matches that of the serialized models
if backend == 'fbgemm':
if backend == "fbgemm":
model.qconfig = torch.quantization.QConfig( # type: ignore[assignment]
activation=torch.quantization.default_observer,
weight=torch.quantization.default_per_channel_weight_observer)
elif backend == 'qnnpack':
weight=torch.quantization.default_per_channel_weight_observer,
)
elif backend == "qnnpack":
model.qconfig = torch.quantization.QConfig( # type: ignore[assignment]
activation=torch.quantization.default_observer,
weight=torch.quantization.default_weight_observer)
activation=torch.quantization.default_observer, weight=torch.quantization.default_weight_observer
)
# TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659
model.fuse_model() # type: ignore[operator]
......
torchvision/models/regnet.py
View file @ 5f0edb97
......@@ -4,11 +4,11 @@
import math
import torch
from collections import OrderedDict
from functools import partial
from typing import Any, Callable, List, Optional, Tuple
import torch
from torch import nn, Tensor
from .._internally_replaced_utils import load_state_dict_from_url
......@@ -16,10 +16,23 @@ from ..ops.misc import ConvNormActivation, SqueezeExcitation
from ._utils import _make_divisible
__all__ = ["RegNet", "regnet_y_400mf", "regnet_y_800mf", "regnet_y_1_6gf",
"regnet_y_3_2gf", "regnet_y_8gf", "regnet_y_16gf", "regnet_y_32gf",
"regnet_x_400mf", "regnet_x_800mf", "regnet_x_1_6gf", "regnet_x_3_2gf",
"regnet_x_8gf", "regnet_x_16gf", "regnet_x_32gf"]
__all__ = [
"RegNet",
"regnet_y_400mf",
"regnet_y_800mf",
"regnet_y_1_6gf",
"regnet_y_3_2gf",
"regnet_y_8gf",
"regnet_y_16gf",
"regnet_y_32gf",
"regnet_x_400mf",
"regnet_x_800mf",
"regnet_x_1_6gf",
"regnet_x_3_2gf",
"regnet_x_8gf",
"regnet_x_16gf",
"regnet_x_32gf",
]
model_urls = {
......@@ -42,8 +55,9 @@ class SimpleStemIN(ConvNormActivation):
norm_layer: Callable[..., nn.Module],
activation_layer: Callable[..., nn.Module],
) -> None:
super().__init__(width_in, width_out, kernel_size=3, stride=2,
norm_layer=norm_layer, activation_layer=activation_layer)
super().__init__(
width_in, width_out, kernel_size=3, stride=2, norm_layer=norm_layer, activation_layer=activation_layer
)
class BottleneckTransform(nn.Sequential):
......@@ -64,10 +78,12 @@ class BottleneckTransform(nn.Sequential):
w_b = int(round(width_out * bottleneck_multiplier))
g = w_b // group_width
layers["a"] = ConvNormActivation(width_in, w_b, kernel_size=1, stride=1,
norm_layer=norm_layer, activation_layer=activation_layer)
layers["b"] = ConvNormActivation(w_b, w_b, kernel_size=3, stride=stride, groups=g,
norm_layer=norm_layer, activation_layer=activation_layer)
layers["a"] = ConvNormActivation(
width_in, w_b, kernel_size=1, stride=1, norm_layer=norm_layer, activation_layer=activation_layer
)
layers["b"] = ConvNormActivation(
w_b, w_b, kernel_size=3, stride=stride, groups=g, norm_layer=norm_layer, activation_layer=activation_layer
)
if se_ratio:
# The SE reduction ratio is defined with respect to the
......@@ -79,8 +95,9 @@ class BottleneckTransform(nn.Sequential):
activation=activation_layer,
)
layers["c"] = ConvNormActivation(w_b, width_out, kernel_size=1, stride=1,
norm_layer=norm_layer, activation_layer=None)
layers["c"] = ConvNormActivation(
w_b, width_out, kernel_size=1, stride=1, norm_layer=norm_layer, activation_layer=None
)
super().__init__(layers)
......@@ -104,8 +121,9 @@ class ResBottleneckBlock(nn.Module):
self.proj = None
should_proj = (width_in != width_out) or (stride != 1)
if should_proj:
self.proj = ConvNormActivation(width_in, width_out, kernel_size=1,
stride=stride, norm_layer=norm_layer, activation_layer=None)
self.proj = ConvNormActivation(
width_in, width_out, kernel_size=1, stride=stride, norm_layer=norm_layer, activation_layer=None
)
self.f = BottleneckTransform(
width_in,
width_out,
......@@ -217,10 +235,7 @@ class BlockParams:
# Compute the block widths. Each stage has one unique block width
widths_cont = torch.arange(depth) * w_a + w_0
block_capacity = torch.round(torch.log(widths_cont / w_0) / math.log(w_m))
block_widths = (
torch.round(torch.divide(w_0 * torch.pow(w_m, block_capacity), QUANT))
* QUANT
).int().tolist()
block_widths = (torch.round(torch.divide(w_0 * torch.pow(w_m, block_capacity), QUANT)) * QUANT).int().tolist()
num_stages = len(set(block_widths))
# Convert to per stage parameters
......@@ -254,14 +269,12 @@ class BlockParams:
)
def _get_expanded_params(self):
return zip(
self.widths, self.strides, self.depths, self.group_widths, self.bottleneck_multipliers
)
return zip(self.widths, self.strides, self.depths, self.group_widths, self.bottleneck_multipliers)
@staticmethod
def _adjust_widths_groups_compatibilty(
stage_widths: List[int], bottleneck_ratios: List[float],
group_widths: List[int]) -> Tuple[List[int], List[int]]:
stage_widths: List[int], bottleneck_ratios: List[float], group_widths: List[int]
) -> Tuple[List[int], List[int]]:
"""
Adjusts the compatibility of widths and groups,
depending on the bottleneck ratio.
......@@ -389,8 +402,7 @@ def regnet_y_400mf(pretrained: bool = False, progress: bool = True, **kwargs: An
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
params = BlockParams.from_init_params(depth=16, w_0=48, w_a=27.89, w_m=2.09,
group_width=8, se_ratio=0.25, **kwargs)
params = BlockParams.from_init_params(depth=16, w_0=48, w_a=27.89, w_m=2.09, group_width=8, se_ratio=0.25, **kwargs)
return _regnet("regnet_y_400mf", params, pretrained, progress, **kwargs)
......@@ -403,8 +415,7 @@ def regnet_y_800mf(pretrained: bool = False, progress: bool = True, **kwargs: An
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
params = BlockParams.from_init_params(depth=14, w_0=56, w_a=38.84, w_m=2.4,
group_width=16, se_ratio=0.25, **kwargs)
params = BlockParams.from_init_params(depth=14, w_0=56, w_a=38.84, w_m=2.4, group_width=16, se_ratio=0.25, **kwargs)
return _regnet("regnet_y_800mf", params, pretrained, progress, **kwargs)
......@@ -417,8 +428,9 @@ def regnet_y_1_6gf(pretrained: bool = False, progress: bool = True, **kwargs: An
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
params = BlockParams.from_init_params(depth=27, w_0=48, w_a=20.71, w_m=2.65,
group_width=24, se_ratio=0.25, **kwargs)
params = BlockParams.from_init_params(
depth=27, w_0=48, w_a=20.71, w_m=2.65, group_width=24, se_ratio=0.25, **kwargs
)
return _regnet("regnet_y_1_6gf", params, pretrained, progress, **kwargs)
......@@ -431,8 +443,9 @@ def regnet_y_3_2gf(pretrained: bool = False, progress: bool = True, **kwargs: An
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
params = BlockParams.from_init_params(depth=21, w_0=80, w_a=42.63, w_m=2.66,
group_width=24, se_ratio=0.25, **kwargs)
params = BlockParams.from_init_params(
depth=21, w_0=80, w_a=42.63, w_m=2.66, group_width=24, se_ratio=0.25, **kwargs
)
return _regnet("regnet_y_3_2gf", params, pretrained, progress, **kwargs)
......@@ -445,8 +458,9 @@ def regnet_y_8gf(pretrained: bool = False, progress: bool = True, **kwargs: Any)
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
params = BlockParams.from_init_params(depth=17, w_0=192, w_a=76.82, w_m=2.19,
group_width=56, se_ratio=0.25, **kwargs)
params = BlockParams.from_init_params(
depth=17, w_0=192, w_a=76.82, w_m=2.19, group_width=56, se_ratio=0.25, **kwargs
)
return _regnet("regnet_y_8gf", params, pretrained, progress, **kwargs)
......@@ -459,8 +473,9 @@ def regnet_y_16gf(pretrained: bool = False, progress: bool = True, **kwargs: Any
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
params = BlockParams.from_init_params(depth=18, w_0=200, w_a=106.23, w_m=2.48,
group_width=112, se_ratio=0.25, **kwargs)
params = BlockParams.from_init_params(
depth=18, w_0=200, w_a=106.23, w_m=2.48, group_width=112, se_ratio=0.25, **kwargs
)
return _regnet("regnet_y_16gf", params, pretrained, progress, **kwargs)
......@@ -473,8 +488,9 @@ def regnet_y_32gf(pretrained: bool = False, progress: bool = True, **kwargs: Any
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
params = BlockParams.from_init_params(depth=20, w_0=232, w_a=115.89, w_m=2.53,
group_width=232, se_ratio=0.25, **kwargs)
params = BlockParams.from_init_params(
depth=20, w_0=232, w_a=115.89, w_m=2.53, group_width=232, se_ratio=0.25, **kwargs
)
return _regnet("regnet_y_32gf", params, pretrained, progress, **kwargs)
......@@ -487,8 +503,7 @@ def regnet_x_400mf(pretrained: bool = False, progress: bool = True, **kwargs: An
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
params = BlockParams.from_init_params(depth=22, w_0=24, w_a=24.48, w_m=2.54,
group_width=16, **kwargs)
params = BlockParams.from_init_params(depth=22, w_0=24, w_a=24.48, w_m=2.54, group_width=16, **kwargs)
return _regnet("regnet_x_400mf", params, pretrained, progress, **kwargs)
......@@ -501,8 +516,7 @@ def regnet_x_800mf(pretrained: bool = False, progress: bool = True, **kwargs: An
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
params = BlockParams.from_init_params(depth=16, w_0=56, w_a=35.73, w_m=2.28,
group_width=16, **kwargs)
params = BlockParams.from_init_params(depth=16, w_0=56, w_a=35.73, w_m=2.28, group_width=16, **kwargs)
return _regnet("regnet_x_800mf", params, pretrained, progress, **kwargs)
......@@ -515,8 +529,7 @@ def regnet_x_1_6gf(pretrained: bool = False, progress: bool = True, **kwargs: An
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
params = BlockParams.from_init_params(depth=18, w_0=80, w_a=34.01, w_m=2.25,
group_width=24, **kwargs)
params = BlockParams.from_init_params(depth=18, w_0=80, w_a=34.01, w_m=2.25, group_width=24, **kwargs)
return _regnet("regnet_x_1_6gf", params, pretrained, progress, **kwargs)
......@@ -529,8 +542,7 @@ def regnet_x_3_2gf(pretrained: bool = False, progress: bool = True, **kwargs: An
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
params = BlockParams.from_init_params(depth=25, w_0=88, w_a=26.31, w_m=2.25,
group_width=48, **kwargs)
params = BlockParams.from_init_params(depth=25, w_0=88, w_a=26.31, w_m=2.25, group_width=48, **kwargs)
return _regnet("regnet_x_3_2gf", params, pretrained, progress, **kwargs)
......@@ -543,8 +555,7 @@ def regnet_x_8gf(pretrained: bool = False, progress: bool = True, **kwargs: Any)
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
params = BlockParams.from_init_params(depth=23, w_0=80, w_a=49.56, w_m=2.88,
group_width=120, **kwargs)
params = BlockParams.from_init_params(depth=23, w_0=80, w_a=49.56, w_m=2.88, group_width=120, **kwargs)
return _regnet("regnet_x_8gf", params, pretrained, progress, **kwargs)
......@@ -557,8 +568,7 @@ def regnet_x_16gf(pretrained: bool = False, progress: bool = True, **kwargs: Any
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
params = BlockParams.from_init_params(depth=22, w_0=216, w_a=55.59, w_m=2.1,
group_width=128, **kwargs)
params = BlockParams.from_init_params(depth=22, w_0=216, w_a=55.59, w_m=2.1, group_width=128, **kwargs)
return _regnet("regnet_x_16gf", params, pretrained, progress, **kwargs)
......@@ -571,8 +581,8 @@ def regnet_x_32gf(pretrained: bool = False, progress: bool = True, **kwargs: Any
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
params = BlockParams.from_init_params(depth=23, w_0=320, w_a=69.86, w_m=2.0,
group_width=168, **kwargs)
params = BlockParams.from_init_params(depth=23, w_0=320, w_a=69.86, w_m=2.0, group_width=168, **kwargs)
return _regnet("regnet_x_32gf", params, pretrained, progress, **kwargs)
# TODO(kazhang): Add RegNetZ_500MF and RegNetZ_4GF
torchvision/models/resnet.py
View file @ 5f0edb97
from typing import Type, Any, Callable, Union, List, Optional
import torch
from torch import Tensor
import torch.nn as nn
from torch import Tensor
from .._internally_replaced_utils import load_state_dict_from_url
from typing import Type, Any, Callable, Union, List, Optional
__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
'resnet152', 'resnext50_32x4d', 'resnext101_32x8d',
'wide_resnet50_2', 'wide_resnet101_2']
__all__ = [
"ResNet",
"resnet18",
"resnet34",
"resnet50",
"resnet101",
"resnet152",
"resnext50_32x4d",
"resnext101_32x8d",
"wide_resnet50_2",
"wide_resnet101_2",
]
model_urls = {
'resnet18': 'https://download.pytorch.org/models/resnet18-f37072fd.pth',
'resnet34': 'https://download.pytorch.org/models/resnet34-b627a593.pth',
'resnet50': 'https://download.pytorch.org/models/resnet50-0676ba61.pth',
'resnet101': 'https://download.pytorch.org/models/resnet101-63fe2227.pth',
'resnet152': 'https://download.pytorch.org/models/resnet152-394f9c45.pth',
'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',
'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth',
"resnet18": "https://download.pytorch.org/models/resnet18-f37072fd.pth",
"resnet34": "https://download.pytorch.org/models/resnet34-b627a593.pth",
"resnet50": "https://download.pytorch.org/models/resnet50-0676ba61.pth",
"resnet101": "https://download.pytorch.org/models/resnet101-63fe2227.pth",
"resnet152": "https://download.pytorch.org/models/resnet152-394f9c45.pth",
"resnext50_32x4d": "https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth",
"resnext101_32x8d": "https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth",
"wide_resnet50_2": "https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth",
"wide_resnet101_2": "https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth",
}
def conv3x3(in_planes: int, out_planes: int, stride: int = 1, groups: int = 1, dilation: int = 1) -> nn.Conv2d:
"""3x3 convolution with padding"""
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
padding=dilation, groups=groups, bias=False, dilation=dilation)
return nn.Conv2d(
in_planes,
out_planes,
kernel_size=3,
stride=stride,
padding=dilation,
groups=groups,
bias=False,
dilation=dilation,
)
def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d:
......@@ -46,13 +65,13 @@ class BasicBlock(nn.Module):
groups: int = 1,
base_width: int = 64,
dilation: int = 1,
norm_layer: Optional[Callable[..., nn.Module]] = None
norm_layer: Optional[Callable[..., nn.Module]] = None,
) -> None:
super(BasicBlock, self).__init__()
if norm_layer is None:
norm_layer = nn.BatchNorm2d
if groups != 1 or base_width != 64:
raise ValueError('BasicBlock only supports groups=1 and base_width=64')
raise ValueError("BasicBlock only supports groups=1 and base_width=64")
if dilation > 1:
raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
# Both self.conv1 and self.downsample layers downsample the input when stride != 1
......@@ -101,12 +120,12 @@ class Bottleneck(nn.Module):
groups: int = 1,
base_width: int = 64,
dilation: int = 1,
norm_layer: Optional[Callable[..., nn.Module]] = None
norm_layer: Optional[Callable[..., nn.Module]] = None,
) -> None:
super(Bottleneck, self).__init__()
if norm_layer is None:
norm_layer = nn.BatchNorm2d
width = int(planes * (base_width / 64.)) * groups
width = int(planes * (base_width / 64.0)) * groups
# Both self.conv2 and self.downsample layers downsample the input when stride != 1
self.conv1 = conv1x1(inplanes, width)
self.bn1 = norm_layer(width)
......@@ -142,7 +161,6 @@ class Bottleneck(nn.Module):
class ResNet(nn.Module):
def __init__(
self,
block: Type[Union[BasicBlock, Bottleneck]],
......@@ -152,7 +170,7 @@ class ResNet(nn.Module):
groups: int = 1,
width_per_group: int = 64,
replace_stride_with_dilation: Optional[List[bool]] = None,
norm_layer: Optional[Callable[..., nn.Module]] = None
norm_layer: Optional[Callable[..., nn.Module]] = None,
) -> None:
super(ResNet, self).__init__()
if norm_layer is None:
......@@ -166,28 +184,26 @@ class ResNet(nn.Module):
# the 2x2 stride with a dilated convolution instead
replace_stride_with_dilation = [False, False, False]
if len(replace_stride_with_dilation) != 3:
raise ValueError("replace_stride_with_dilation should be None "
"or a 3-element tuple, got {}".format(replace_stride_with_dilation))
raise ValueError(
"replace_stride_with_dilation should be None "
"or a 3-element tuple, got {}".format(replace_stride_with_dilation)
)
self.groups = groups
self.base_width = width_per_group
self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3,
bias=False)
self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)
self.bn1 = norm_layer(self.inplanes)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(block, 128, layers[1], stride=2,
dilate=replace_stride_with_dilation[0])
self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
dilate=replace_stride_with_dilation[1])
self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
dilate=replace_stride_with_dilation[2])
self.layer2 = self._make_layer(block, 128, layers[1], stride=2, dilate=replace_stride_with_dilation[0])
self.layer3 = self._make_layer(block, 256, layers[2], stride=2, dilate=replace_stride_with_dilation[1])
self.layer4 = self._make_layer(block, 512, layers[3], stride=2, dilate=replace_stride_with_dilation[2])
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(512 * block.expansion, num_classes)
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
......@@ -202,8 +218,14 @@ class ResNet(nn.Module):
elif isinstance(m, BasicBlock):
nn.init.constant_(m.bn2.weight, 0) # type: ignore[arg-type]
def _make_layer(self, block: Type[Union[BasicBlock, Bottleneck]], planes: int, blocks: int,
stride: int = 1, dilate: bool = False) -> nn.Sequential:
def _make_layer(
self,
block: Type[Union[BasicBlock, Bottleneck]],
planes: int,
blocks: int,
stride: int = 1,
dilate: bool = False,
) -> nn.Sequential:
norm_layer = self._norm_layer
downsample = None
previous_dilation = self.dilation
......@@ -217,13 +239,23 @@ class ResNet(nn.Module):
)
layers = []
layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
self.base_width, previous_dilation, norm_layer))
layers.append(
block(
self.inplanes, planes, stride, downsample, self.groups, self.base_width, previous_dilation, norm_layer
)
)
self.inplanes = planes * block.expansion
for _ in range(1, blocks):
layers.append(block(self.inplanes, planes, groups=self.groups,
base_width=self.base_width, dilation=self.dilation,
norm_layer=norm_layer))
layers.append(
block(
self.inplanes,
planes,
groups=self.groups,
base_width=self.base_width,
dilation=self.dilation,
norm_layer=norm_layer,
)
)
return nn.Sequential(*layers)
......@@ -255,12 +287,11 @@ def _resnet(
layers: List[int],
pretrained: bool,
progress: bool,
**kwargs: Any
**kwargs: Any,
) -> ResNet:
model = ResNet(block, layers, **kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls[arch],
progress=progress)
state_dict = load_state_dict_from_url(model_urls[arch], progress=progress)
model.load_state_dict(state_dict)
return model
......@@ -273,8 +304,7 @@ def resnet18(pretrained: bool = False, progress: bool = True, **kwargs: Any) ->
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress,
**kwargs)
return _resnet("resnet18", BasicBlock, [2, 2, 2, 2], pretrained, progress, **kwargs)
def resnet34(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
......@@ -285,8 +315,7 @@ def resnet34(pretrained: bool = False, progress: bool = True, **kwargs: Any) ->
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress,
**kwargs)
return _resnet("resnet34", BasicBlock, [3, 4, 6, 3], pretrained, progress, **kwargs)
def resnet50(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
......@@ -297,8 +326,7 @@ def resnet50(pretrained: bool = False, progress: bool = True, **kwargs: Any) ->
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress,
**kwargs)
return _resnet("resnet50", Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs)
def resnet101(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
......@@ -309,8 +337,7 @@ def resnet101(pretrained: bool = False, progress: bool = True, **kwargs: Any) ->
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress,
**kwargs)
return _resnet("resnet101", Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs)
def resnet152(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
......@@ -321,8 +348,7 @@ def resnet152(pretrained: bool = False, progress: bool = True, **kwargs: Any) ->
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress,
**kwargs)
return _resnet("resnet152", Bottleneck, [3, 8, 36, 3], pretrained, progress, **kwargs)
def resnext50_32x4d(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
......@@ -333,10 +359,9 @@ def resnext50_32x4d(pretrained: bool = False, progress: bool = True, **kwargs: A
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
kwargs['groups'] = 32
kwargs['width_per_group'] = 4
return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3],
pretrained, progress, **kwargs)
kwargs["groups"] = 32
kwargs["width_per_group"] = 4
return _resnet("resnext50_32x4d", Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs)
def resnext101_32x8d(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
......@@ -347,10 +372,9 @@ def resnext101_32x8d(pretrained: bool = False, progress: bool = True, **kwargs:
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
kwargs['groups'] = 32
kwargs['width_per_group'] = 8
return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3],
pretrained, progress, **kwargs)
kwargs["groups"] = 32
kwargs["width_per_group"] = 8
return _resnet("resnext101_32x8d", Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs)
def wide_resnet50_2(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
......@@ -366,9 +390,8 @@ def wide_resnet50_2(pretrained: bool = False, progress: bool = True, **kwargs: A
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
kwargs['width_per_group'] = 64 * 2
return _resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3],
pretrained, progress, **kwargs)
kwargs["width_per_group"] = 64 * 2
return _resnet("wide_resnet50_2", Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs)
def wide_resnet101_2(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
......@@ -384,6 +407,5 @@ def wide_resnet101_2(pretrained: bool = False, progress: bool = True, **kwargs:
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
kwargs['width_per_group'] = 64 * 2
return _resnet('wide_resnet101_2', Bottleneck, [3, 4, 23, 3],
pretrained, progress, **kwargs)
kwargs["width_per_group"] = 64 * 2
return _resnet("wide_resnet101_2", Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs)
torchvision/models/segmentation/_utils.py
View file @ 5f0edb97
......@@ -6,14 +6,9 @@ from torch.nn import functional as F
class _SimpleSegmentationModel(nn.Module):
__constants__ = ['aux_classifier']
def __init__(
self,
backbone: nn.Module,
classifier: nn.Module,
aux_classifier: Optional[nn.Module] = None
) -> None:
__constants__ = ["aux_classifier"]
def __init__(self, backbone: nn.Module, classifier: nn.Module, aux_classifier: Optional[nn.Module] = None) -> None:
super(_SimpleSegmentationModel, self).__init__()
self.backbone = backbone
self.classifier = classifier
......@@ -27,13 +22,13 @@ class _SimpleSegmentationModel(nn.Module):
result = OrderedDict()
x = features["out"]
x = self.classifier(x)
x = F.interpolate(x, size=input_shape, mode='bilinear', align_corners=False)
x = F.interpolate(x, size=input_shape, mode="bilinear", align_corners=False)
result["out"] = x
if self.aux_classifier is not None:
x = features["aux"]
x = self.aux_classifier(x)
x = F.interpolate(x, size=input_shape, mode='bilinear', align_corners=False)
x = F.interpolate(x, size=input_shape, mode="bilinear", align_corners=False)
result["aux"] = x
return result
torchvision/models/segmentation/deeplabv3.py
View file @ 5f0edb97
from typing import List
import torch
from torch import nn
from torch.nn import functional as F
from typing import List
from ._utils import _SimpleSegmentationModel
......@@ -24,6 +25,7 @@ class DeepLabV3(_SimpleSegmentationModel):
the backbone and returns a dense prediction.
aux_classifier (nn.Module, optional): auxiliary classifier used during training
"""
pass
......@@ -34,7 +36,7 @@ class DeepLabHead(nn.Sequential):
nn.Conv2d(256, 256, 3, padding=1, bias=False),
nn.BatchNorm2d(256),
nn.ReLU(),
nn.Conv2d(256, num_classes, 1)
nn.Conv2d(256, num_classes, 1),
)
......@@ -43,7 +45,7 @@ class ASPPConv(nn.Sequential):
modules = [
nn.Conv2d(in_channels, out_channels, 3, padding=dilation, dilation=dilation, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU()
nn.ReLU(),
]
super(ASPPConv, self).__init__(*modules)
......@@ -54,23 +56,23 @@ class ASPPPooling(nn.Sequential):
nn.AdaptiveAvgPool2d(1),
nn.Conv2d(in_channels, out_channels, 1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU())
nn.ReLU(),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
size = x.shape[-2:]
for mod in self:
x = mod(x)
return F.interpolate(x, size=size, mode='bilinear', align_corners=False)
return F.interpolate(x, size=size, mode="bilinear", align_corners=False)
class ASPP(nn.Module):
def __init__(self, in_channels: int, atrous_rates: List[int], out_channels: int = 256) -> None:
super(ASPP, self).__init__()
modules = []
modules.append(nn.Sequential(
nn.Conv2d(in_channels, out_channels, 1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU()))
modules.append(
nn.Sequential(nn.Conv2d(in_channels, out_channels, 1, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU())
)
rates = tuple(atrous_rates)
for rate in rates:
......@@ -84,7 +86,8 @@ class ASPP(nn.Module):
nn.Conv2d(len(self.convs) * out_channels, out_channels, 1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(),
nn.Dropout(0.5))
nn.Dropout(0.5),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
_res = []
......
torchvision/models/segmentation/fcn.py
View file @ 5f0edb97
......@@ -19,6 +19,7 @@ class FCN(_SimpleSegmentationModel):
the backbone and returns a dense prediction.
aux_classifier (nn.Module, optional): auxiliary classifier used during training
"""
pass
......@@ -30,7 +31,7 @@ class FCNHead(nn.Sequential):
nn.BatchNorm2d(inter_channels),
nn.ReLU(),
nn.Dropout(0.1),
nn.Conv2d(inter_channels, channels, 1)
nn.Conv2d(inter_channels, channels, 1),
]
super(FCNHead, self).__init__(*layers)
torchvision/models/segmentation/lraspp.py
View file @ 5f0edb97
from collections import OrderedDict
from typing import Dict
from torch import nn, Tensor
from torch.nn import functional as F
from typing import Dict
__all__ = ["LRASPP"]
......@@ -25,12 +25,7 @@ class LRASPP(nn.Module):
"""
def __init__(
self,
backbone: nn.Module,
low_channels: int,
high_channels: int,
num_classes: int,
inter_channels: int = 128
self, backbone: nn.Module, low_channels: int, high_channels: int, num_classes: int, inter_channels: int = 128
) -> None:
super().__init__()
self.backbone = backbone
......@@ -39,7 +34,7 @@ class LRASPP(nn.Module):
def forward(self, input: Tensor) -> Dict[str, Tensor]:
features = self.backbone(input)
out = self.classifier(features)
out = F.interpolate(out, size=input.shape[-2:], mode='bilinear', align_corners=False)
out = F.interpolate(out, size=input.shape[-2:], mode="bilinear", align_corners=False)
result = OrderedDict()
result["out"] = out
......@@ -48,19 +43,12 @@ class LRASPP(nn.Module):
class LRASPPHead(nn.Module):
def __init__(
self,
low_channels: int,
high_channels: int,
num_classes: int,
inter_channels: int
) -> None:
def __init__(self, low_channels: int, high_channels: int, num_classes: int, inter_channels: int) -> None:
super().__init__()
self.cbr = nn.Sequential(
nn.Conv2d(high_channels, inter_channels, 1, bias=False),
nn.BatchNorm2d(inter_channels),
nn.ReLU(inplace=True)
nn.ReLU(inplace=True),
)
self.scale = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
......@@ -77,6 +65,6 @@ class LRASPPHead(nn.Module):
x = self.cbr(high)
s = self.scale(high)
x = x * s
x = F.interpolate(x, size=low.shape[-2:], mode='bilinear', align_corners=False)
x = F.interpolate(x, size=low.shape[-2:], mode="bilinear", align_corners=False)
return self.low_classifier(low) + self.high_classifier(x)
torchvision/models/segmentation/segmentation.py
View file @ 5f0edb97
from torch import nn
from typing import Any, Optional
from .._utils import IntermediateLayerGetter
from torch import nn
from ..._internally_replaced_utils import load_state_dict_from_url
from .. import mobilenetv3
from .. import resnet
from .._utils import IntermediateLayerGetter
from .deeplabv3 import DeepLabHead, DeepLabV3
from .fcn import FCN, FCNHead
from .lraspp import LRASPP
__all__ = ['fcn_resnet50', 'fcn_resnet101', 'deeplabv3_resnet50', 'deeplabv3_resnet101',
'deeplabv3_mobilenet_v3_large', 'lraspp_mobilenet_v3_large']
__all__ = [
"fcn_resnet50",
"fcn_resnet101",
"deeplabv3_resnet50",
"deeplabv3_resnet101",
"deeplabv3_mobilenet_v3_large",
"lraspp_mobilenet_v3_large",
]
model_urls = {
'fcn_resnet50_coco': 'https://download.pytorch.org/models/fcn_resnet50_coco-1167a1af.pth',
'fcn_resnet101_coco': 'https://download.pytorch.org/models/fcn_resnet101_coco-7ecb50ca.pth',
'deeplabv3_resnet50_coco': 'https://download.pytorch.org/models/deeplabv3_resnet50_coco-cd0a2569.pth',
'deeplabv3_resnet101_coco': 'https://download.pytorch.org/models/deeplabv3_resnet101_coco-586e9e4e.pth',
'deeplabv3_mobilenet_v3_large_coco':
'https://download.pytorch.org/models/deeplabv3_mobilenet_v3_large-fc3c493d.pth',
'lraspp_mobilenet_v3_large_coco': 'https://download.pytorch.org/models/lraspp_mobilenet_v3_large-d234d4ea.pth',
"fcn_resnet50_coco": "https://download.pytorch.org/models/fcn_resnet50_coco-1167a1af.pth",
"fcn_resnet101_coco": "https://download.pytorch.org/models/fcn_resnet101_coco-7ecb50ca.pth",
"deeplabv3_resnet50_coco": "https://download.pytorch.org/models/deeplabv3_resnet50_coco-cd0a2569.pth",
"deeplabv3_resnet101_coco": "https://download.pytorch.org/models/deeplabv3_resnet101_coco-586e9e4e.pth",
"deeplabv3_mobilenet_v3_large_coco": "https://download.pytorch.org/models/deeplabv3_mobilenet_v3_large-fc3c493d.pth",
"lraspp_mobilenet_v3_large_coco": "https://download.pytorch.org/models/lraspp_mobilenet_v3_large-d234d4ea.pth",
}
def _segm_model(
name: str,
backbone_name: str,
num_classes: int,
aux: Optional[bool],
pretrained_backbone: bool = True
name: str, backbone_name: str, num_classes: int, aux: Optional[bool], pretrained_backbone: bool = True
) -> nn.Module:
if 'resnet' in backbone_name:
if "resnet" in backbone_name:
backbone = resnet.__dict__[backbone_name](
pretrained=pretrained_backbone,
replace_stride_with_dilation=[False, True, True])
out_layer = 'layer4'
pretrained=pretrained_backbone, replace_stride_with_dilation=[False, True, True]
)
out_layer = "layer4"
out_inplanes = 2048
aux_layer = 'layer3'
aux_layer = "layer3"
aux_inplanes = 1024
elif 'mobilenet_v3' in backbone_name:
elif "mobilenet_v3" in backbone_name:
backbone = mobilenetv3.__dict__[backbone_name](pretrained=pretrained_backbone, dilated=True).features
# Gather the indices of blocks which are strided. These are the locations of C1, ..., Cn-1 blocks.
......@@ -52,11 +55,11 @@ def _segm_model(
aux_layer = str(aux_pos)
aux_inplanes = backbone[aux_pos].out_channels
else:
raise NotImplementedError('backbone {} is not supported as of now'.format(backbone_name))
raise NotImplementedError("backbone {} is not supported as of now".format(backbone_name))
return_layers = {out_layer: 'out'}
return_layers = {out_layer: "out"}
if aux:
return_layers[aux_layer] = 'aux'
return_layers[aux_layer] = "aux"
backbone = IntermediateLayerGetter(backbone, return_layers=return_layers)
aux_classifier = None
......@@ -64,8 +67,8 @@ def _segm_model(
aux_classifier = FCNHead(aux_inplanes, num_classes)
model_map = {
'deeplabv3': (DeepLabHead, DeepLabV3),
'fcn': (FCNHead, FCN),
"deeplabv3": (DeepLabHead, DeepLabV3),
"fcn": (FCNHead, FCN),
}
classifier = model_map[name][0](out_inplanes, num_classes)
base_model = model_map[name][1]
......@@ -81,7 +84,7 @@ def _load_model(
progress: bool,
num_classes: int,
aux_loss: Optional[bool],
**kwargs: Any
**kwargs: Any,
) -> nn.Module:
if pretrained:
aux_loss = True
......@@ -93,10 +96,10 @@ def _load_model(
def _load_weights(model: nn.Module, arch_type: str, backbone: str, progress: bool) -> None:
arch = arch_type + '_' + backbone + '_coco'
arch = arch_type + "_" + backbone + "_coco"
model_url = model_urls.get(arch, None)
if model_url is None:
raise NotImplementedError('pretrained {} is not supported as of now'.format(arch))
raise NotImplementedError("pretrained {} is not supported as of now".format(arch))
else:
state_dict = load_state_dict_from_url(model_url, progress=progress)
model.load_state_dict(state_dict)
......@@ -113,7 +116,7 @@ def _segm_lraspp_mobilenetv3(backbone_name: str, num_classes: int, pretrained_ba
low_channels = backbone[low_pos].out_channels
high_channels = backbone[high_pos].out_channels
backbone = IntermediateLayerGetter(backbone, return_layers={str(low_pos): 'low', str(high_pos): 'high'})
backbone = IntermediateLayerGetter(backbone, return_layers={str(low_pos): "low", str(high_pos): "high"})
model = LRASPP(backbone, low_channels, high_channels, num_classes)
return model
......@@ -124,7 +127,7 @@ def fcn_resnet50(
progress: bool = True,
num_classes: int = 21,
aux_loss: Optional[bool] = None,
**kwargs: Any
**kwargs: Any,
) -> nn.Module:
"""Constructs a Fully-Convolutional Network model with a ResNet-50 backbone.
......@@ -135,7 +138,7 @@ def fcn_resnet50(
num_classes (int): number of output classes of the model (including the background)
aux_loss (bool): If True, it uses an auxiliary loss
"""
return _load_model('fcn', 'resnet50', pretrained, progress, num_classes, aux_loss, **kwargs)
return _load_model("fcn", "resnet50", pretrained, progress, num_classes, aux_loss, **kwargs)
def fcn_resnet101(
......@@ -143,7 +146,7 @@ def fcn_resnet101(
progress: bool = True,
num_classes: int = 21,
aux_loss: Optional[bool] = None,
**kwargs: Any
**kwargs: Any,
) -> nn.Module:
"""Constructs a Fully-Convolutional Network model with a ResNet-101 backbone.
......@@ -154,7 +157,7 @@ def fcn_resnet101(
num_classes (int): number of output classes of the model (including the background)
aux_loss (bool): If True, it uses an auxiliary loss
"""
return _load_model('fcn', 'resnet101', pretrained, progress, num_classes, aux_loss, **kwargs)
return _load_model("fcn", "resnet101", pretrained, progress, num_classes, aux_loss, **kwargs)
def deeplabv3_resnet50(
......@@ -162,7 +165,7 @@ def deeplabv3_resnet50(
progress: bool = True,
num_classes: int = 21,
aux_loss: Optional[bool] = None,
**kwargs: Any
**kwargs: Any,
) -> nn.Module:
"""Constructs a DeepLabV3 model with a ResNet-50 backbone.
......@@ -173,7 +176,7 @@ def deeplabv3_resnet50(
num_classes (int): number of output classes of the model (including the background)
aux_loss (bool): If True, it uses an auxiliary loss
"""
return _load_model('deeplabv3', 'resnet50', pretrained, progress, num_classes, aux_loss, **kwargs)
return _load_model("deeplabv3", "resnet50", pretrained, progress, num_classes, aux_loss, **kwargs)
def deeplabv3_resnet101(
......@@ -181,7 +184,7 @@ def deeplabv3_resnet101(
progress: bool = True,
num_classes: int = 21,
aux_loss: Optional[bool] = None,
**kwargs: Any
**kwargs: Any,
) -> nn.Module:
"""Constructs a DeepLabV3 model with a ResNet-101 backbone.
......@@ -192,7 +195,7 @@ def deeplabv3_resnet101(
num_classes (int): The number of classes
aux_loss (bool): If True, include an auxiliary classifier
"""
return _load_model('deeplabv3', 'resnet101', pretrained, progress, num_classes, aux_loss, **kwargs)
return _load_model("deeplabv3", "resnet101", pretrained, progress, num_classes, aux_loss, **kwargs)
def deeplabv3_mobilenet_v3_large(
......@@ -200,7 +203,7 @@ def deeplabv3_mobilenet_v3_large(
progress: bool = True,
num_classes: int = 21,
aux_loss: Optional[bool] = None,
**kwargs: Any
**kwargs: Any,
) -> nn.Module:
"""Constructs a DeepLabV3 model with a MobileNetV3-Large backbone.
......@@ -211,14 +214,11 @@ def deeplabv3_mobilenet_v3_large(
num_classes (int): number of output classes of the model (including the background)
aux_loss (bool): If True, it uses an auxiliary loss
"""
return _load_model('deeplabv3', 'mobilenet_v3_large', pretrained, progress, num_classes, aux_loss, **kwargs)
return _load_model("deeplabv3", "mobilenet_v3_large", pretrained, progress, num_classes, aux_loss, **kwargs)
def lraspp_mobilenet_v3_large(
pretrained: bool = False,
progress: bool = True,
num_classes: int = 21,
**kwargs: Any
pretrained: bool = False, progress: bool = True, num_classes: int = 21, **kwargs: Any
) -> nn.Module:
"""Constructs a Lite R-ASPP Network model with a MobileNetV3-Large backbone.
......@@ -229,14 +229,14 @@ def lraspp_mobilenet_v3_large(
num_classes (int): number of output classes of the model (including the background)
"""
if kwargs.pop("aux_loss", False):
raise NotImplementedError('This model does not use auxiliary loss')
raise NotImplementedError("This model does not use auxiliary loss")
backbone_name = 'mobilenet_v3_large'
backbone_name = "mobilenet_v3_large"
if pretrained:
kwargs["pretrained_backbone"] = False
model = _segm_lraspp_mobilenetv3(backbone_name, num_classes, **kwargs)
if pretrained:
_load_weights(model, 'lraspp', backbone_name, progress)
_load_weights(model, "lraspp", backbone_name, progress)
return model
torchvision/models/shufflenetv2.py
View file @ 5f0edb97
from typing import Callable, Any, List
import torch
from torch import Tensor
import torch.nn as nn
from torch import Tensor
from .._internally_replaced_utils import load_state_dict_from_url
from typing import Callable, Any, List
__all__ = [
'ShuffleNetV2', 'shufflenet_v2_x0_5', 'shufflenet_v2_x1_0',
'shufflenet_v2_x1_5', 'shufflenet_v2_x2_0'
]
__all__ = ["ShuffleNetV2", "shufflenet_v2_x0_5", "shufflenet_v2_x1_0", "shufflenet_v2_x1_5", "shufflenet_v2_x2_0"]
model_urls = {
'shufflenetv2_x0.5': 'https://download.pytorch.org/models/shufflenetv2_x0.5-f707e7126e.pth',
'shufflenetv2_x1.0': 'https://download.pytorch.org/models/shufflenetv2_x1-5666bf0f80.pth',
'shufflenetv2_x1.5': None,
'shufflenetv2_x2.0': None,
"shufflenetv2_x0.5": "https://download.pytorch.org/models/shufflenetv2_x0.5-f707e7126e.pth",
"shufflenetv2_x1.0": "https://download.pytorch.org/models/shufflenetv2_x1-5666bf0f80.pth",
"shufflenetv2_x1.5": None,
"shufflenetv2_x2.0": None,
}
......@@ -23,8 +22,7 @@ def channel_shuffle(x: Tensor, groups: int) -> Tensor:
channels_per_group = num_channels // groups
# reshape
x = x.view(batchsize, groups,
channels_per_group, height, width)
x = x.view(batchsize, groups, channels_per_group, height, width)
x = torch.transpose(x, 1, 2).contiguous()
......@@ -35,16 +33,11 @@ def channel_shuffle(x: Tensor, groups: int) -> Tensor:
class InvertedResidual(nn.Module):
def __init__(
self,
inp: int,
oup: int,
stride: int
) -> None:
def __init__(self, inp: int, oup: int, stride: int) -> None:
super(InvertedResidual, self).__init__()
if not (1 <= stride <= 3):
raise ValueError('illegal stride value')
raise ValueError("illegal stride value")
self.stride = stride
branch_features = oup // 2
......@@ -62,8 +55,14 @@ class InvertedResidual(nn.Module):
self.branch1 = nn.Sequential()
self.branch2 = nn.Sequential(
nn.Conv2d(inp if (self.stride > 1) else branch_features,
branch_features, kernel_size=1, stride=1, padding=0, bias=False),
nn.Conv2d(
inp if (self.stride > 1) else branch_features,
branch_features,
kernel_size=1,
stride=1,
padding=0,
bias=False,
),
nn.BatchNorm2d(branch_features),
nn.ReLU(inplace=True),
self.depthwise_conv(branch_features, branch_features, kernel_size=3, stride=self.stride, padding=1),
......@@ -75,12 +74,7 @@ class InvertedResidual(nn.Module):
@staticmethod
def depthwise_conv(
i: int,
o: int,
kernel_size: int,
stride: int = 1,
padding: int = 0,
bias: bool = False
i: int, o: int, kernel_size: int, stride: int = 1, padding: int = 0, bias: bool = False
) -> nn.Conv2d:
return nn.Conv2d(i, o, kernel_size, stride, padding, bias=bias, groups=i)
......@@ -102,14 +96,14 @@ class ShuffleNetV2(nn.Module):
stages_repeats: List[int],
stages_out_channels: List[int],
num_classes: int = 1000,
inverted_residual: Callable[..., nn.Module] = InvertedResidual
inverted_residual: Callable[..., nn.Module] = InvertedResidual,
) -> None:
super(ShuffleNetV2, self).__init__()
if len(stages_repeats) != 3:
raise ValueError('expected stages_repeats as list of 3 positive ints')
raise ValueError("expected stages_repeats as list of 3 positive ints")
if len(stages_out_channels) != 5:
raise ValueError('expected stages_out_channels as list of 5 positive ints')
raise ValueError("expected stages_out_channels as list of 5 positive ints")
self._stage_out_channels = stages_out_channels
input_channels = 3
......@@ -127,9 +121,8 @@ class ShuffleNetV2(nn.Module):
self.stage2: nn.Sequential
self.stage3: nn.Sequential
self.stage4: nn.Sequential
stage_names = ['stage{}'.format(i) for i in [2, 3, 4]]
for name, repeats, output_channels in zip(
stage_names, stages_repeats, self._stage_out_channels[1:]):
stage_names = ["stage{}".format(i) for i in [2, 3, 4]]
for name, repeats, output_channels in zip(stage_names, stages_repeats, self._stage_out_channels[1:]):
seq = [inverted_residual(input_channels, output_channels, 2)]
for i in range(repeats - 1):
seq.append(inverted_residual(output_channels, output_channels, 1))
......@@ -167,7 +160,7 @@ def _shufflenetv2(arch: str, pretrained: bool, progress: bool, *args: Any, **kwa
if pretrained:
model_url = model_urls[arch]
if model_url is None:
raise NotImplementedError('pretrained {} is not supported as of now'.format(arch))
raise NotImplementedError("pretrained {} is not supported as of now".format(arch))
else:
state_dict = load_state_dict_from_url(model_url, progress=progress)
model.load_state_dict(state_dict)
......@@ -185,8 +178,7 @@ def shufflenet_v2_x0_5(pretrained: bool = False, progress: bool = True, **kwargs
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _shufflenetv2('shufflenetv2_x0.5', pretrained, progress,
[4, 8, 4], [24, 48, 96, 192, 1024], **kwargs)
return _shufflenetv2("shufflenetv2_x0.5", pretrained, progress, [4, 8, 4], [24, 48, 96, 192, 1024], **kwargs)
def shufflenet_v2_x1_0(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ShuffleNetV2:
......@@ -199,8 +191,7 @@ def shufflenet_v2_x1_0(pretrained: bool = False, progress: bool = True, **kwargs
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _shufflenetv2('shufflenetv2_x1.0', pretrained, progress,
[4, 8, 4], [24, 116, 232, 464, 1024], **kwargs)
return _shufflenetv2("shufflenetv2_x1.0", pretrained, progress, [4, 8, 4], [24, 116, 232, 464, 1024], **kwargs)
def shufflenet_v2_x1_5(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ShuffleNetV2:
......@@ -213,8 +204,7 @@ def shufflenet_v2_x1_5(pretrained: bool = False, progress: bool = True, **kwargs
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _shufflenetv2('shufflenetv2_x1.5', pretrained, progress,
[4, 8, 4], [24, 176, 352, 704, 1024], **kwargs)
return _shufflenetv2("shufflenetv2_x1.5", pretrained, progress, [4, 8, 4], [24, 176, 352, 704, 1024], **kwargs)
def shufflenet_v2_x2_0(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ShuffleNetV2:
......@@ -227,5 +217,4 @@ def shufflenet_v2_x2_0(pretrained: bool = False, progress: bool = True, **kwargs
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _shufflenetv2('shufflenetv2_x2.0', pretrained, progress,
[4, 8, 4], [24, 244, 488, 976, 2048], **kwargs)
return _shufflenetv2("shufflenetv2_x2.0", pretrained, progress, [4, 8, 4], [24, 244, 488, 976, 2048], **kwargs)
torchvision/models/squeezenet.py
View file @ 5f0edb97
from typing import Any
import torch
import torch.nn as nn
import torch.nn.init as init
from .._internally_replaced_utils import load_state_dict_from_url
from typing import Any
__all__ = ['SqueezeNet', 'squeezenet1_0', 'squeezenet1_1']
__all__ = ["SqueezeNet", "squeezenet1_0", "squeezenet1_1"]
model_urls = {
'squeezenet1_0': 'https://download.pytorch.org/models/squeezenet1_0-b66bff10.pth',
'squeezenet1_1': 'https://download.pytorch.org/models/squeezenet1_1-b8a52dc0.pth',
"squeezenet1_0": "https://download.pytorch.org/models/squeezenet1_0-b66bff10.pth",
"squeezenet1_1": "https://download.pytorch.org/models/squeezenet1_1-b8a52dc0.pth",
}
class Fire(nn.Module):
def __init__(
self,
inplanes: int,
squeeze_planes: int,
expand1x1_planes: int,
expand3x3_planes: int
) -> None:
def __init__(self, inplanes: int, squeeze_planes: int, expand1x1_planes: int, expand3x3_planes: int) -> None:
super(Fire, self).__init__()
self.inplanes = inplanes
self.squeeze = nn.Conv2d(inplanes, squeeze_planes, kernel_size=1)
self.squeeze_activation = nn.ReLU(inplace=True)
self.expand1x1 = nn.Conv2d(squeeze_planes, expand1x1_planes,
kernel_size=1)
self.expand1x1 = nn.Conv2d(squeeze_planes, expand1x1_planes, kernel_size=1)
self.expand1x1_activation = nn.ReLU(inplace=True)
self.expand3x3 = nn.Conv2d(squeeze_planes, expand3x3_planes,
kernel_size=3, padding=1)
self.expand3x3 = nn.Conv2d(squeeze_planes, expand3x3_planes, kernel_size=3, padding=1)
self.expand3x3_activation = nn.ReLU(inplace=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.squeeze_activation(self.squeeze(x))
return torch.cat([
self.expand1x1_activation(self.expand1x1(x)),
self.expand3x3_activation(self.expand3x3(x))
], 1)
return torch.cat(
[self.expand1x1_activation(self.expand1x1(x)), self.expand3x3_activation(self.expand3x3(x))], 1
)
class SqueezeNet(nn.Module):
def __init__(
self,
version: str = '1_0',
num_classes: int = 1000
) -> None:
def __init__(self, version: str = "1_0", num_classes: int = 1000) -> None:
super(SqueezeNet, self).__init__()
self.num_classes = num_classes
if version == '1_0':
if version == "1_0":
self.features = nn.Sequential(
nn.Conv2d(3, 96, kernel_size=7, stride=2),
nn.ReLU(inplace=True),
......@@ -65,7 +52,7 @@ class SqueezeNet(nn.Module):
nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
Fire(512, 64, 256, 256),
)
elif version == '1_1':
elif version == "1_1":
self.features = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=3, stride=2),
nn.ReLU(inplace=True),
......@@ -85,16 +72,12 @@ class SqueezeNet(nn.Module):
# FIXME: Is this needed? SqueezeNet should only be called from the
# FIXME: squeezenet1_x() functions
# FIXME: This checking is not done for the other models
raise ValueError("Unsupported SqueezeNet version {version}:"
"1_0 or 1_1 expected".format(version=version))
raise ValueError("Unsupported SqueezeNet version {version}:" "1_0 or 1_1 expected".format(version=version))
# Final convolution is initialized differently from the rest
final_conv = nn.Conv2d(512, self.num_classes, kernel_size=1)
self.classifier = nn.Sequential(
nn.Dropout(p=0.5),
final_conv,
nn.ReLU(inplace=True),
nn.AdaptiveAvgPool2d((1, 1))
nn.Dropout(p=0.5), final_conv, nn.ReLU(inplace=True), nn.AdaptiveAvgPool2d((1, 1))
)
for m in self.modules():
......@@ -115,9 +98,8 @@ class SqueezeNet(nn.Module):
def _squeezenet(version: str, pretrained: bool, progress: bool, **kwargs: Any) -> SqueezeNet:
model = SqueezeNet(version, **kwargs)
if pretrained:
arch = 'squeezenet' + version
state_dict = load_state_dict_from_url(model_urls[arch],
progress=progress)
arch = "squeezenet" + version
state_dict = load_state_dict_from_url(model_urls[arch], progress=progress)
model.load_state_dict(state_dict)
return model
......@@ -132,7 +114,7 @@ def squeezenet1_0(pretrained: bool = False, progress: bool = True, **kwargs: Any
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _squeezenet('1_0', pretrained, progress, **kwargs)
return _squeezenet("1_0", pretrained, progress, **kwargs)
def squeezenet1_1(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> SqueezeNet:
......@@ -146,4 +128,4 @@ def squeezenet1_1(pretrained: bool = False, progress: bool = True, **kwargs: Any
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _squeezenet('1_1', pretrained, progress, **kwargs)
return _squeezenet("1_1", pretrained, progress, **kwargs)
torchvision/models/vgg.py
View file @ 5f0edb97
from typing import Union, List, Dict, Any, cast
import torch
import torch.nn as nn
from .._internally_replaced_utils import load_state_dict_from_url
from typing import Union, List, Dict, Any, cast
__all__ = [
'VGG', 'vgg11', 'vgg11_bn', 'vgg13', 'vgg13_bn', 'vgg16', 'vgg16_bn',
'vgg19_bn', 'vgg19',
"VGG",
"vgg11",
"vgg11_bn",
"vgg13",
"vgg13_bn",
"vgg16",
"vgg16_bn",
"vgg19_bn",
"vgg19",
]
model_urls = {
'vgg11': 'https://download.pytorch.org/models/vgg11-8a719046.pth',
'vgg13': 'https://download.pytorch.org/models/vgg13-19584684.pth',
'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth',
'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth',
'vgg11_bn': 'https://download.pytorch.org/models/vgg11_bn-6002323d.pth',
'vgg13_bn': 'https://download.pytorch.org/models/vgg13_bn-abd245e5.pth',
'vgg16_bn': 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth',
'vgg19_bn': 'https://download.pytorch.org/models/vgg19_bn-c79401a0.pth',
"vgg11": "https://download.pytorch.org/models/vgg11-8a719046.pth",
"vgg13": "https://download.pytorch.org/models/vgg13-19584684.pth",
"vgg16": "https://download.pytorch.org/models/vgg16-397923af.pth",
"vgg19": "https://download.pytorch.org/models/vgg19-dcbb9e9d.pth",
"vgg11_bn": "https://download.pytorch.org/models/vgg11_bn-6002323d.pth",
"vgg13_bn": "https://download.pytorch.org/models/vgg13_bn-abd245e5.pth",
"vgg16_bn": "https://download.pytorch.org/models/vgg16_bn-6c64b313.pth",
"vgg19_bn": "https://download.pytorch.org/models/vgg19_bn-c79401a0.pth",
}
class VGG(nn.Module):
def __init__(
self,
features: nn.Module,
num_classes: int = 1000,
init_weights: bool = True
) -> None:
def __init__(self, features: nn.Module, num_classes: int = 1000, init_weights: bool = True) -> None:
super(VGG, self).__init__()
self.features = features
self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
......@@ -55,7 +58,7 @@ class VGG(nn.Module):
def _initialize_weights(self) -> None:
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
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):
......@@ -70,7 +73,7 @@ def make_layers(cfg: List[Union[str, int]], batch_norm: bool = False) -> nn.Sequ
layers: List[nn.Module] = []
in_channels = 3
for v in cfg:
if v == 'M':
if v == "M":
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
else:
v = cast(int, v)
......@@ -84,20 +87,19 @@ def make_layers(cfg: List[Union[str, int]], batch_norm: bool = False) -> nn.Sequ
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'],
"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(arch: str, cfg: str, batch_norm: bool, pretrained: bool, progress: bool, **kwargs: Any) -> VGG:
if pretrained:
kwargs['init_weights'] = False
kwargs["init_weights"] = False
model = VGG(make_layers(cfgs[cfg], batch_norm=batch_norm), **kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls[arch],
progress=progress)
state_dict = load_state_dict_from_url(model_urls[arch], progress=progress)
model.load_state_dict(state_dict)
return model
......@@ -111,7 +113,7 @@ def vgg11(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VGG
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _vgg('vgg11', 'A', False, pretrained, progress, **kwargs)
return _vgg("vgg11", "A", False, pretrained, progress, **kwargs)
def vgg11_bn(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VGG:
......@@ -123,7 +125,7 @@ def vgg11_bn(pretrained: bool = False, progress: bool = True, **kwargs: Any) ->
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _vgg('vgg11_bn', 'A', True, pretrained, progress, **kwargs)
return _vgg("vgg11_bn", "A", True, pretrained, progress, **kwargs)
def vgg13(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VGG:
......@@ -135,7 +137,7 @@ def vgg13(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VGG
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _vgg('vgg13', 'B', False, pretrained, progress, **kwargs)
return _vgg("vgg13", "B", False, pretrained, progress, **kwargs)
def vgg13_bn(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VGG:
......@@ -147,7 +149,7 @@ def vgg13_bn(pretrained: bool = False, progress: bool = True, **kwargs: Any) ->
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _vgg('vgg13_bn', 'B', True, pretrained, progress, **kwargs)
return _vgg("vgg13_bn", "B", True, pretrained, progress, **kwargs)
def vgg16(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VGG:
......@@ -159,7 +161,7 @@ def vgg16(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VGG
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _vgg('vgg16', 'D', False, pretrained, progress, **kwargs)
return _vgg("vgg16", "D", False, pretrained, progress, **kwargs)
def vgg16_bn(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VGG:
......@@ -171,7 +173,7 @@ def vgg16_bn(pretrained: bool = False, progress: bool = True, **kwargs: Any) ->
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _vgg('vgg16_bn', 'D', True, pretrained, progress, **kwargs)
return _vgg("vgg16_bn", "D", True, pretrained, progress, **kwargs)
def vgg19(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VGG:
......@@ -183,7 +185,7 @@ def vgg19(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VGG
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _vgg('vgg19', 'E', False, pretrained, progress, **kwargs)
return _vgg("vgg19", "E", False, pretrained, progress, **kwargs)
def vgg19_bn(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VGG:
......@@ -195,4 +197,4 @@ def vgg19_bn(pretrained: bool = False, progress: bool = True, **kwargs: Any) ->
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _vgg('vgg19_bn', 'E', True, pretrained, progress, **kwargs)
return _vgg("vgg19_bn", "E", True, pretrained, progress, **kwargs)
torchvision/models/video/resnet.py
View file @ 5f0edb97
from torch import Tensor
import torch.nn as nn
from typing import Tuple, Optional, Callable, List, Type, Any, Union
import torch.nn as nn
from torch import Tensor
from ..._internally_replaced_utils import load_state_dict_from_url
__all__ = ['r3d_18', 'mc3_18', 'r2plus1d_18']
__all__ = ["r3d_18", "mc3_18", "r2plus1d_18"]
model_urls = {
'r3d_18': 'https://download.pytorch.org/models/r3d_18-b3b3357e.pth',
'mc3_18': 'https://download.pytorch.org/models/mc3_18-a90a0ba3.pth',
'r2plus1d_18': 'https://download.pytorch.org/models/r2plus1d_18-91a641e6.pth',
"r3d_18": "https://download.pytorch.org/models/r3d_18-b3b3357e.pth",
"mc3_18": "https://download.pytorch.org/models/mc3_18-a90a0ba3.pth",
"r2plus1d_18": "https://download.pytorch.org/models/r2plus1d_18-91a641e6.pth",
}
class Conv3DSimple(nn.Conv3d):
def __init__(
self,
in_planes: int,
out_planes: int,
midplanes: Optional[int] = None,
stride: int = 1,
padding: int = 1
self, in_planes: int, out_planes: int, midplanes: Optional[int] = None, stride: int = 1, padding: int = 1
) -> None:
super(Conv3DSimple, self).__init__(
......@@ -30,7 +26,8 @@ class Conv3DSimple(nn.Conv3d):
kernel_size=(3, 3, 3),
stride=stride,
padding=padding,
bias=False)
bias=False,
)
@staticmethod
def get_downsample_stride(stride: int) -> Tuple[int, int, int]:
......@@ -38,24 +35,22 @@ class Conv3DSimple(nn.Conv3d):
class Conv2Plus1D(nn.Sequential):
def __init__(
self,
in_planes: int,
out_planes: int,
midplanes: int,
stride: int = 1,
padding: int = 1
) -> None:
def __init__(self, in_planes: int, out_planes: int, midplanes: int, stride: int = 1, padding: int = 1) -> None:
super(Conv2Plus1D, self).__init__(
nn.Conv3d(in_planes, midplanes, kernel_size=(1, 3, 3),
stride=(1, stride, stride), padding=(0, padding, padding),
bias=False),
nn.Conv3d(
in_planes,
midplanes,
kernel_size=(1, 3, 3),
stride=(1, stride, stride),
padding=(0, padding, padding),
bias=False,
),
nn.BatchNorm3d(midplanes),
nn.ReLU(inplace=True),
nn.Conv3d(midplanes, out_planes, kernel_size=(3, 1, 1),
stride=(stride, 1, 1), padding=(padding, 0, 0),
bias=False))
nn.Conv3d(
midplanes, out_planes, kernel_size=(3, 1, 1), stride=(stride, 1, 1), padding=(padding, 0, 0), bias=False
),
)
@staticmethod
def get_downsample_stride(stride: int) -> Tuple[int, int, int]:
......@@ -63,14 +58,8 @@ class Conv2Plus1D(nn.Sequential):
class Conv3DNoTemporal(nn.Conv3d):
def __init__(
self,
in_planes: int,
out_planes: int,
midplanes: Optional[int] = None,
stride: int = 1,
padding: int = 1
self, in_planes: int, out_planes: int, midplanes: Optional[int] = None, stride: int = 1, padding: int = 1
) -> None:
super(Conv3DNoTemporal, self).__init__(
......@@ -79,7 +68,8 @@ class Conv3DNoTemporal(nn.Conv3d):
kernel_size=(1, 3, 3),
stride=(1, stride, stride),
padding=(0, padding, padding),
bias=False)
bias=False,
)
@staticmethod
def get_downsample_stride(stride: int) -> Tuple[int, int, int]:
......@@ -102,14 +92,9 @@ class BasicBlock(nn.Module):
super(BasicBlock, self).__init__()
self.conv1 = nn.Sequential(
conv_builder(inplanes, planes, midplanes, stride),
nn.BatchNorm3d(planes),
nn.ReLU(inplace=True)
)
self.conv2 = nn.Sequential(
conv_builder(planes, planes, midplanes),
nn.BatchNorm3d(planes)
conv_builder(inplanes, planes, midplanes, stride), nn.BatchNorm3d(planes), nn.ReLU(inplace=True)
)
self.conv2 = nn.Sequential(conv_builder(planes, planes, midplanes), nn.BatchNorm3d(planes))
self.relu = nn.ReLU(inplace=True)
self.downsample = downsample
self.stride = stride
......@@ -145,21 +130,17 @@ class Bottleneck(nn.Module):
# 1x1x1
self.conv1 = nn.Sequential(
nn.Conv3d(inplanes, planes, kernel_size=1, bias=False),
nn.BatchNorm3d(planes),
nn.ReLU(inplace=True)
nn.Conv3d(inplanes, planes, kernel_size=1, bias=False), nn.BatchNorm3d(planes), nn.ReLU(inplace=True)
)
# Second kernel
self.conv2 = nn.Sequential(
conv_builder(planes, planes, midplanes, stride),
nn.BatchNorm3d(planes),
nn.ReLU(inplace=True)
conv_builder(planes, planes, midplanes, stride), nn.BatchNorm3d(planes), nn.ReLU(inplace=True)
)
# 1x1x1
self.conv3 = nn.Sequential(
nn.Conv3d(planes, planes * self.expansion, kernel_size=1, bias=False),
nn.BatchNorm3d(planes * self.expansion)
nn.BatchNorm3d(planes * self.expansion),
)
self.relu = nn.ReLU(inplace=True)
self.downsample = downsample
......@@ -182,35 +163,31 @@ class Bottleneck(nn.Module):
class BasicStem(nn.Sequential):
"""The default conv-batchnorm-relu stem
"""
"""The default conv-batchnorm-relu stem"""
def __init__(self) -> None:
super(BasicStem, self).__init__(
nn.Conv3d(3, 64, kernel_size=(3, 7, 7), stride=(1, 2, 2),
padding=(1, 3, 3), bias=False),
nn.Conv3d(3, 64, kernel_size=(3, 7, 7), stride=(1, 2, 2), padding=(1, 3, 3), bias=False),
nn.BatchNorm3d(64),
nn.ReLU(inplace=True))
nn.ReLU(inplace=True),
)
class R2Plus1dStem(nn.Sequential):
"""R(2+1)D stem is different than the default one as it uses separated 3D convolution
"""
"""R(2+1)D stem is different than the default one as it uses separated 3D convolution"""
def __init__(self) -> None:
super(R2Plus1dStem, self).__init__(
nn.Conv3d(3, 45, kernel_size=(1, 7, 7),
stride=(1, 2, 2), padding=(0, 3, 3),
bias=False),
nn.Conv3d(3, 45, kernel_size=(1, 7, 7), stride=(1, 2, 2), padding=(0, 3, 3), bias=False),
nn.BatchNorm3d(45),
nn.ReLU(inplace=True),
nn.Conv3d(45, 64, kernel_size=(3, 1, 1),
stride=(1, 1, 1), padding=(1, 0, 0),
bias=False),
nn.Conv3d(45, 64, kernel_size=(3, 1, 1), stride=(1, 1, 1), padding=(1, 0, 0), bias=False),
nn.BatchNorm3d(64),
nn.ReLU(inplace=True))
nn.ReLU(inplace=True),
)
class VideoResNet(nn.Module):
def __init__(
self,
block: Type[Union[BasicBlock, Bottleneck]],
......@@ -273,16 +250,15 @@ class VideoResNet(nn.Module):
conv_builder: Type[Union[Conv3DSimple, Conv3DNoTemporal, Conv2Plus1D]],
planes: int,
blocks: int,
stride: int = 1
stride: int = 1,
) -> nn.Sequential:
downsample = None
if stride != 1 or self.inplanes != planes * block.expansion:
ds_stride = conv_builder.get_downsample_stride(stride)
downsample = nn.Sequential(
nn.Conv3d(self.inplanes, planes * block.expansion,
kernel_size=1, stride=ds_stride, bias=False),
nn.BatchNorm3d(planes * block.expansion)
nn.Conv3d(self.inplanes, planes * block.expansion, kernel_size=1, stride=ds_stride, bias=False),
nn.BatchNorm3d(planes * block.expansion),
)
layers = []
layers.append(block(self.inplanes, planes, conv_builder, stride, downsample))
......@@ -296,8 +272,7 @@ class VideoResNet(nn.Module):
def _initialize_weights(self) -> None:
for m in self.modules():
if isinstance(m, nn.Conv3d):
nn.init.kaiming_normal_(m.weight, mode='fan_out',
nonlinearity='relu')
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.BatchNorm3d):
......@@ -312,8 +287,7 @@ def _video_resnet(arch: str, pretrained: bool = False, progress: bool = True, **
model = VideoResNet(**kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls[arch],
progress=progress)
state_dict = load_state_dict_from_url(model_urls[arch], progress=progress)
model.load_state_dict(state_dict)
return model
......@@ -330,12 +304,16 @@ def r3d_18(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> Vi
nn.Module: R3D-18 network
"""
return _video_resnet('r3d_18',
pretrained, progress,
block=BasicBlock,
conv_makers=[Conv3DSimple] * 4,
layers=[2, 2, 2, 2],
stem=BasicStem, **kwargs)
return _video_resnet(
"r3d_18",
pretrained,
progress,
block=BasicBlock,
conv_makers=[Conv3DSimple] * 4,
layers=[2, 2, 2, 2],
stem=BasicStem,
**kwargs,
)
def mc3_18(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VideoResNet:
......@@ -349,12 +327,16 @@ def mc3_18(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> Vi
Returns:
nn.Module: MC3 Network definition
"""
return _video_resnet('mc3_18',
pretrained, progress,
block=BasicBlock,
conv_makers=[Conv3DSimple] + [Conv3DNoTemporal] * 3, # type: ignore[list-item]
layers=[2, 2, 2, 2],
stem=BasicStem, **kwargs)
return _video_resnet(
"mc3_18",
pretrained,
progress,
block=BasicBlock,
conv_makers=[Conv3DSimple] + [Conv3DNoTemporal] * 3, # type: ignore[list-item]
layers=[2, 2, 2, 2],
stem=BasicStem,
**kwargs,
)
def r2plus1d_18(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VideoResNet:
......@@ -368,9 +350,13 @@ def r2plus1d_18(pretrained: bool = False, progress: bool = True, **kwargs: Any)
Returns:
nn.Module: R(2+1)D-18 network
"""
return _video_resnet('r2plus1d_18',
pretrained, progress,
block=BasicBlock,
conv_makers=[Conv2Plus1D] * 4,
layers=[2, 2, 2, 2],
stem=R2Plus1dStem, **kwargs)
return _video_resnet(
"r2plus1d_18",
pretrained,
progress,
block=BasicBlock,
conv_makers=[Conv2Plus1D] * 4,
layers=[2, 2, 2, 2],
stem=R2Plus1dStem,
**kwargs,
)
torchvision/ops/__init__.py
View file @ 5f0edb97
from .boxes import nms, batched_nms, remove_small_boxes, clip_boxes_to_image, box_area, box_iou, generalized_box_iou, \
masks_to_boxes
from ._register_onnx_ops import _register_custom_op
from .boxes import (
nms,
batched_nms,
remove_small_boxes,
clip_boxes_to_image,
box_area,
box_iou,
generalized_box_iou,
masks_to_boxes,
)
from .boxes import box_convert
from .deform_conv import deform_conv2d, DeformConv2d
from .roi_align import roi_align, RoIAlign
from .roi_pool import roi_pool, RoIPool
from .ps_roi_align import ps_roi_align, PSRoIAlign
from .ps_roi_pool import ps_roi_pool, PSRoIPool
from .poolers import MultiScaleRoIAlign
from .feature_pyramid_network import FeaturePyramidNetwork
from .focal_loss import sigmoid_focal_loss
from .poolers import MultiScaleRoIAlign
from .ps_roi_align import ps_roi_align, PSRoIAlign
from .ps_roi_pool import ps_roi_pool, PSRoIPool
from .roi_align import roi_align, RoIAlign
from .roi_pool import roi_pool, RoIPool
from .stochastic_depth import stochastic_depth, StochasticDepth
from ._register_onnx_ops import _register_custom_op
_register_custom_op()
__all__ = [
'deform_conv2d', 'DeformConv2d', 'nms', 'batched_nms', 'remove_small_boxes',
'clip_boxes_to_image', 'box_convert',
'box_area', 'box_iou', 'generalized_box_iou', 'roi_align', 'RoIAlign', 'roi_pool',
'RoIPool', 'ps_roi_align', 'PSRoIAlign', 'ps_roi_pool',
'PSRoIPool', 'MultiScaleRoIAlign', 'FeaturePyramidNetwork',
'sigmoid_focal_loss', 'stochastic_depth', 'StochasticDepth'
"deform_conv2d",
"DeformConv2d",
"nms",
"batched_nms",
"remove_small_boxes",
"clip_boxes_to_image",
"box_convert",
"box_area",
"box_iou",
"generalized_box_iou",
"roi_align",
"RoIAlign",
"roi_pool",
"RoIPool",
"ps_roi_align",
"PSRoIAlign",
"ps_roi_pool",
"PSRoIPool",
"MultiScaleRoIAlign",
"FeaturePyramidNetwork",
"sigmoid_focal_loss",
"stochastic_depth",
"StochasticDepth",
]
torchvision/ops/_register_onnx_ops.py
View file @ 5f0edb97
import sys
import torch
import warnings
import torch
_onnx_opset_version = 11
def _register_custom_op():
from torch.onnx.symbolic_helper import parse_args, scalar_type_to_onnx, scalar_type_to_pytorch_type, \
cast_pytorch_to_onnx
from torch.onnx.symbolic_opset9 import _cast_Long
from torch.onnx.symbolic_helper import (
parse_args,
scalar_type_to_onnx,
scalar_type_to_pytorch_type,
cast_pytorch_to_onnx,
)
from torch.onnx.symbolic_opset11 import select, squeeze, unsqueeze
from torch.onnx.symbolic_opset9 import _cast_Long
@parse_args('v', 'v', 'f')
@parse_args("v", "v", "f")
def symbolic_multi_label_nms(g, boxes, scores, iou_threshold):
boxes = unsqueeze(g, boxes, 0)
scores = unsqueeze(g, unsqueeze(g, scores, 0), 0)
max_output_per_class = g.op('Constant', value_t=torch.tensor([sys.maxsize], dtype=torch.long))
iou_threshold = g.op('Constant', value_t=torch.tensor([iou_threshold], dtype=torch.float))
nms_out = g.op('NonMaxSuppression', boxes, scores, max_output_per_class, iou_threshold)
return squeeze(g, select(g, nms_out, 1, g.op('Constant', value_t=torch.tensor([2], dtype=torch.long))), 1)
max_output_per_class = g.op("Constant", value_t=torch.tensor([sys.maxsize], dtype=torch.long))
iou_threshold = g.op("Constant", value_t=torch.tensor([iou_threshold], dtype=torch.float))
nms_out = g.op("NonMaxSuppression", boxes, scores, max_output_per_class, iou_threshold)
return squeeze(g, select(g, nms_out, 1, g.op("Constant", value_t=torch.tensor([2], dtype=torch.long))), 1)
@parse_args('v', 'v', 'f', 'i', 'i', 'i', 'i')
@parse_args("v", "v", "f", "i", "i", "i", "i")
def roi_align(g, input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio, aligned):
batch_indices = _cast_Long(g, squeeze(g, select(g, rois, 1, g.op('Constant',
value_t=torch.tensor([0], dtype=torch.long))), 1), False)
rois = select(g, rois, 1, g.op('Constant', value_t=torch.tensor([1, 2, 3, 4], dtype=torch.long)))
batch_indices = _cast_Long(
g, squeeze(g, select(g, rois, 1, g.op("Constant", value_t=torch.tensor([0], dtype=torch.long))), 1), False
)
rois = select(g, rois, 1, g.op("Constant", value_t=torch.tensor([1, 2, 3, 4], dtype=torch.long)))
if aligned:
warnings.warn("ONNX export of ROIAlign with aligned=True does not match PyTorch when using malformed boxes,"
" ONNX forces ROIs to be 1x1 or larger.")
warnings.warn(
"ONNX export of ROIAlign with aligned=True does not match PyTorch when using malformed boxes,"
" ONNX forces ROIs to be 1x1 or larger."
)
scale = torch.tensor(0.5 / spatial_scale).to(dtype=torch.float)
rois = g.op("Sub", rois, scale)
# ONNX doesn't support negative sampling_ratio
if sampling_ratio < 0:
warnings.warn("ONNX doesn't support negative sampling ratio,"
"therefore is is set to 0 in order to be exported.")
warnings.warn(
"ONNX doesn't support negative sampling ratio," "therefore is is set to 0 in order to be exported."
)
sampling_ratio = 0
return g.op('RoiAlign', input, rois, batch_indices, spatial_scale_f=spatial_scale,
output_height_i=pooled_height, output_width_i=pooled_width, sampling_ratio_i=sampling_ratio)
return g.op(
"RoiAlign",
input,
rois,
batch_indices,
spatial_scale_f=spatial_scale,
output_height_i=pooled_height,
output_width_i=pooled_width,
sampling_ratio_i=sampling_ratio,
)
@parse_args('v', 'v', 'f', 'i', 'i')
@parse_args("v", "v", "f", "i", "i")
def roi_pool(g, input, rois, spatial_scale, pooled_height, pooled_width):
roi_pool = g.op('MaxRoiPool', input, rois,
pooled_shape_i=(pooled_height, pooled_width), spatial_scale_f=spatial_scale)
roi_pool = g.op(
"MaxRoiPool", input, rois, pooled_shape_i=(pooled_height, pooled_width), spatial_scale_f=spatial_scale
)
return roi_pool, None
from torch.onnx import register_custom_op_symbolic
register_custom_op_symbolic('torchvision::nms', symbolic_multi_label_nms, _onnx_opset_version)
register_custom_op_symbolic('torchvision::roi_align', roi_align, _onnx_opset_version)
register_custom_op_symbolic('torchvision::roi_pool', roi_pool, _onnx_opset_version)
register_custom_op_symbolic("torchvision::nms", symbolic_multi_label_nms, _onnx_opset_version)
register_custom_op_symbolic("torchvision::roi_align", roi_align, _onnx_opset_version)
register_custom_op_symbolic("torchvision::roi_pool", roi_pool, _onnx_opset_version)
torchvision/ops/_utils.py
View file @ 5f0edb97
from typing import List, Union
import torch
from torch import Tensor
from typing import List, Union
def _cat(tensors: List[Tensor], dim: int = 0) -> Tensor:
......@@ -27,10 +28,11 @@ def convert_boxes_to_roi_format(boxes: List[Tensor]) -> Tensor:
def check_roi_boxes_shape(boxes: Union[Tensor, List[Tensor]]):
if isinstance(boxes, (list, tuple)):
for _tensor in boxes:
assert _tensor.size(1) == 4, \
'The shape of the tensor in the boxes list is not correct as List[Tensor[L, 4]]'
assert (
_tensor.size(1) == 4
), "The shape of the tensor in the boxes list is not correct as List[Tensor[L, 4]]"
elif isinstance(boxes, torch.Tensor):
assert boxes.size(1) == 5, 'The boxes tensor shape is not correct as Tensor[K, 5]'
assert boxes.size(1) == 5, "The boxes tensor shape is not correct as Tensor[K, 5]"
else:
assert False, 'boxes is expected to be a Tensor[L, 5] or a List[Tensor[K, 4]]'
assert False, "boxes is expected to be a Tensor[L, 5] or a List[Tensor[K, 4]]"
return
torchvision/ops/boxes.py
View file @ 5f0edb97
import torch
from torch import Tensor
from typing import Tuple
from ._box_convert import _box_cxcywh_to_xyxy, _box_xyxy_to_cxcywh, _box_xywh_to_xyxy, _box_xyxy_to_xywh
import torch
import torchvision
from torch import Tensor
from torchvision.extension import _assert_has_ops
from ._box_convert import _box_cxcywh_to_xyxy, _box_xyxy_to_cxcywh, _box_xywh_to_xyxy, _box_xyxy_to_xywh
def nms(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor:
"""
......@@ -183,13 +185,13 @@ def box_convert(boxes: Tensor, in_fmt: str, out_fmt: str) -> Tensor:
if in_fmt == out_fmt:
return boxes.clone()
if in_fmt != 'xyxy' and out_fmt != 'xyxy':
if in_fmt != "xyxy" and out_fmt != "xyxy":
# convert to xyxy and change in_fmt xyxy
if in_fmt == "xywh":
boxes = _box_xywh_to_xyxy(boxes)
elif in_fmt == "cxcywh":
boxes = _box_cxcywh_to_xyxy(boxes)
in_fmt = 'xyxy'
in_fmt = "xyxy"
if in_fmt == "xyxy":
if out_fmt == "xywh":
......
torchvision/ops/deform_conv.py
View file @ 5f0edb97
import math
from typing import Optional, Tuple
import torch
from torch import nn, Tensor
from torch.nn import init
from torch.nn.parameter import Parameter
from torch.nn.modules.utils import _pair
from typing import Optional, Tuple
from torch.nn.parameter import Parameter
from torchvision.extension import _assert_has_ops
......@@ -84,7 +84,9 @@ def deform_conv2d(
"the shape of the offset tensor at dimension 1 is not valid. It should "
"be a multiple of 2 * weight.size[2] * weight.size[3].\n"
"Got offset.shape[1]={}, while 2 * weight.size[2] * weight.size[3]={}".format(
offset.shape[1], 2 * weights_h * weights_w))
offset.shape[1], 2 * weights_h * weights_w
)
)
return torch.ops.torchvision.deform_conv2d(
input,
......@@ -92,12 +94,16 @@ def deform_conv2d(
offset,
mask,
bias,
stride_h, stride_w,
pad_h, pad_w,
dil_h, dil_w,
stride_h,
stride_w,
pad_h,
pad_w,
dil_h,
dil_w,
n_weight_grps,
n_offset_grps,
use_mask,)
use_mask,
)
class DeformConv2d(nn.Module):
......@@ -119,9 +125,9 @@ class DeformConv2d(nn.Module):
super(DeformConv2d, self).__init__()
if in_channels % groups != 0:
raise ValueError('in_channels must be divisible by groups')
raise ValueError("in_channels must be divisible by groups")
if out_channels % groups != 0:
raise ValueError('out_channels must be divisible by groups')
raise ValueError("out_channels must be divisible by groups")
self.in_channels = in_channels
self.out_channels = out_channels
......@@ -131,13 +137,14 @@ class DeformConv2d(nn.Module):
self.dilation = _pair(dilation)
self.groups = groups
self.weight = Parameter(torch.empty(out_channels, in_channels // groups,
self.kernel_size[0], self.kernel_size[1]))
self.weight = Parameter(
torch.empty(out_channels, in_channels // groups, self.kernel_size[0], self.kernel_size[1])
)
if bias:
self.bias = Parameter(torch.empty(out_channels))
else:
self.register_parameter('bias', None)
self.register_parameter("bias", None)
self.reset_parameters()
......@@ -160,18 +167,26 @@ class DeformConv2d(nn.Module):
out_height, out_width]): masks to be applied for each position in the
convolution kernel.
"""
return deform_conv2d(input, offset, self.weight, self.bias, stride=self.stride,
padding=self.padding, dilation=self.dilation, mask=mask)
return deform_conv2d(
input,
offset,
self.weight,
self.bias,
stride=self.stride,
padding=self.padding,
dilation=self.dilation,
mask=mask,
)
def __repr__(self) -> str:
s = self.__class__.__name__ + '('
s += '{in_channels}'
s += ', {out_channels}'
s += ', kernel_size={kernel_size}'
s += ', stride={stride}'
s += ', padding={padding}' if self.padding != (0, 0) else ''
s += ', dilation={dilation}' if self.dilation != (1, 1) else ''
s += ', groups={groups}' if self.groups != 1 else ''
s += ', bias=False' if self.bias is None else ''
s += ')'
s = self.__class__.__name__ + "("
s += "{in_channels}"
s += ", {out_channels}"
s += ", kernel_size={kernel_size}"
s += ", stride={stride}"
s += ", padding={padding}" if self.padding != (0, 0) else ""
s += ", dilation={dilation}" if self.dilation != (1, 1) else ""
s += ", groups={groups}" if self.groups != 1 else ""
s += ", bias=False" if self.bias is None else ""
s += ")"
return s.format(**self.__dict__)
torchvision/ops/feature_pyramid_network.py
View file @ 5f0edb97
from collections import OrderedDict
from typing import Tuple, List, Dict, Optional
import torch.nn.functional as F
from torch import nn, Tensor
from typing import Tuple, List, Dict, Optional
class ExtraFPNBlock(nn.Module):
"""
......@@ -21,6 +20,7 @@ class ExtraFPNBlock(nn.Module):
of the FPN
names (List[str]): the extended set of names for the results
"""
def forward(
self,
results: List[Tensor],
......@@ -67,6 +67,7 @@ class FeaturePyramidNetwork(nn.Module):
>>> ('feat3', torch.Size([1, 5, 8, 8]))]
"""
def __init__(
self,
in_channels_list: List[int],
......@@ -165,6 +166,7 @@ class LastLevelMaxPool(ExtraFPNBlock):
"""
Applies a max_pool2d on top of the last feature map
"""
def forward(
self,
x: List[Tensor],
......@@ -180,6 +182,7 @@ class LastLevelP6P7(ExtraFPNBlock):
"""
This module is used in RetinaNet to generate extra layers, P6 and P7.
"""
def __init__(self, in_channels: int, out_channels: int):
super(LastLevelP6P7, self).__init__()
self.p6 = nn.Conv2d(in_channels, out_channels, 3, 2, 1)
......
torchvision/ops/focal_loss.py
View file @ 5f0edb97
......@@ -31,9 +31,7 @@ def sigmoid_focal_loss(
Loss tensor with the reduction option applied.
"""
p = torch.sigmoid(inputs)
ce_loss = F.binary_cross_entropy_with_logits(
inputs, targets, reduction="none"
)
ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none")
p_t = p * targets + (1 - p) * (1 - targets)
loss = ce_loss * ((1 - p_t) ** gamma)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment