backbones.py 2.61 KB
Newer Older
Abdullah Rashwan's avatar
Abdullah Rashwan committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# Lint as: python3
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Backbones configurations."""
Abdullah Rashwan's avatar
Abdullah Rashwan committed
17
from typing import Optional, List
Abdullah Rashwan's avatar
Abdullah Rashwan committed
18
19
20
21
22
23
24
25
26
27
28

# Import libraries
import dataclasses

from official.modeling import hyperparams


@dataclasses.dataclass
class ResNet(hyperparams.Config):
  """ResNet config."""
  model_id: int = 50
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
29
30
  stem_type: str = 'v0'
  se_ratio: float = 0.0
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
31
  stochastic_depth_drop_rate: float = 0.0
Abdullah Rashwan's avatar
Abdullah Rashwan committed
32
33


Abdullah Rashwan's avatar
Abdullah Rashwan committed
34
35
36
37
38
@dataclasses.dataclass
class DilatedResNet(hyperparams.Config):
  """DilatedResNet config."""
  model_id: int = 50
  output_stride: int = 16
Abdullah Rashwan's avatar
Abdullah Rashwan committed
39
40
41
  multigrid: Optional[List[int]] = None
  stem_type: str = 'v0'
  last_stage_repeats: int = 1
Abdullah Rashwan's avatar
Abdullah Rashwan committed
42
43


Abdullah Rashwan's avatar
Abdullah Rashwan committed
44
45
46
47
48
@dataclasses.dataclass
class EfficientNet(hyperparams.Config):
  """EfficientNet config."""
  model_id: str = 'b0'
  se_ratio: float = 0.0
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
49
  stochastic_depth_drop_rate: float = 0.0
Abdullah Rashwan's avatar
Abdullah Rashwan committed
50
51


52
53
54
55
@dataclasses.dataclass
class MobileNet(hyperparams.Config):
  """Mobilenet config."""
  model_id: str = 'MobileNetV2'
56
  filter_size_scale: float = 1.0
57
58
59
  stochastic_depth_drop_rate: float = 0.0


Abdullah Rashwan's avatar
Abdullah Rashwan committed
60
61
62
63
@dataclasses.dataclass
class SpineNet(hyperparams.Config):
  """SpineNet config."""
  model_id: str = '49'
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
64
  stochastic_depth_drop_rate: float = 0.0
Abdullah Rashwan's avatar
Abdullah Rashwan committed
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80


@dataclasses.dataclass
class RevNet(hyperparams.Config):
  """RevNet config."""
  # Specifies the depth of RevNet.
  model_id: int = 56


@dataclasses.dataclass
class Backbone(hyperparams.OneOfConfig):
  """Configuration for backbones.

  Attributes:
    type: 'str', type of backbone be used, one the of fields below.
    resnet: resnet backbone config.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
81
    dilated_resnet: dilated resnet backbone for semantic segmentation config.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
82
83
84
    revnet: revnet backbone config.
    efficientnet: efficientnet backbone config.
    spinenet: spinenet backbone config.
85
    mobilenet: mobilenet backbone config.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
86
87
88
  """
  type: Optional[str] = None
  resnet: ResNet = ResNet()
Abdullah Rashwan's avatar
Abdullah Rashwan committed
89
  dilated_resnet: DilatedResNet = DilatedResNet()
Abdullah Rashwan's avatar
Abdullah Rashwan committed
90
91
92
  revnet: RevNet = RevNet()
  efficientnet: EfficientNet = EfficientNet()
  spinenet: SpineNet = SpineNet()
93
  mobilenet: MobileNet = MobileNet()