(
self,
growth_rate: int = 32,
block_config: Tuple[int, int, int, int] = (6, 12, 24, 16),
num_init_features: int = 64,
bn_size: int = 4,
drop_rate: float = 0,
num_classes: int = 1000,
memory_efficient: bool = False
)
| 152 | """ |
| 153 | |
| 154 | def __init__( |
| 155 | self, |
| 156 | growth_rate: int = 32, |
| 157 | block_config: Tuple[int, int, int, int] = (6, 12, 24, 16), |
| 158 | num_init_features: int = 64, |
| 159 | bn_size: int = 4, |
| 160 | drop_rate: float = 0, |
| 161 | num_classes: int = 1000, |
| 162 | memory_efficient: bool = False |
| 163 | ) -> None: |
| 164 | |
| 165 | super(DenseNet, self).__init__() |
| 166 | |
| 167 | # First convolution |
| 168 | self.features = nn.Sequential(OrderedDict([ |
| 169 | ('conv0', nn.Conv2d(3, num_init_features, kernel_size=7, stride=2, |
| 170 | padding=3, bias=False)), |
| 171 | ('norm0', nn.BatchNorm2d(num_init_features)), |
| 172 | ('relu0', nn.ReLU(inplace=True)), |
| 173 | ('pool0', nn.MaxPool2d(kernel_size=3, stride=2, padding=1)), |
| 174 | ])) |
| 175 | |
| 176 | # Each denseblock |
| 177 | num_features = num_init_features |
| 178 | for i, num_layers in enumerate(block_config): |
| 179 | block = _DenseBlock( |
| 180 | num_layers=num_layers, |
| 181 | num_input_features=num_features, |
| 182 | bn_size=bn_size, |
| 183 | growth_rate=growth_rate, |
| 184 | drop_rate=drop_rate, |
| 185 | memory_efficient=memory_efficient |
| 186 | ) |
| 187 | self.features.add_module('denseblock%d' % (i + 1), block) |
| 188 | num_features = num_features + num_layers * growth_rate |
| 189 | if i != len(block_config) - 1: |
| 190 | trans = _Transition(num_input_features=num_features, |
| 191 | num_output_features=num_features // 2) |
| 192 | self.features.add_module('transition%d' % (i + 1), trans) |
| 193 | num_features = num_features // 2 |
| 194 | |
| 195 | # Final batch norm |
| 196 | self.features.add_module('norm5', nn.BatchNorm2d(num_features)) |
| 197 | |
| 198 | # Linear layer |
| 199 | self.classifier = nn.Linear(num_features, num_classes) |
| 200 | |
| 201 | # Official init from torch repo. |
| 202 | for m in self.modules(): |
| 203 | if isinstance(m, nn.Conv2d): |
| 204 | nn.init.kaiming_normal_(m.weight) |
| 205 | elif isinstance(m, nn.BatchNorm2d): |
| 206 | nn.init.constant_(m.weight, 1) |
| 207 | nn.init.constant_(m.bias, 0) |
| 208 | elif isinstance(m, nn.Linear): |
| 209 | nn.init.constant_(m.bias, 0) |
| 210 | |
| 211 | def forward(self, x: Tensor) -> Tensor: |
no test coverage detected