| 17 | |
| 18 | class _DenseLayer(nn.Module): |
| 19 | def __init__( |
| 20 | self, |
| 21 | num_input_features: int, |
| 22 | growth_rate: int, |
| 23 | bn_size: int, |
| 24 | drop_rate: float, |
| 25 | memory_efficient: bool = False |
| 26 | ) -> None: |
| 27 | super(_DenseLayer, self).__init__() |
| 28 | self.norm1: nn.BatchNorm2d |
| 29 | self.add_module('norm1', nn.BatchNorm2d(num_input_features)) |
| 30 | self.relu1: nn.ReLU |
| 31 | self.add_module('relu1', nn.ReLU(inplace=True)) |
| 32 | self.conv1: nn.Conv2d |
| 33 | self.add_module('conv1', nn.Conv2d(num_input_features, bn_size * |
| 34 | growth_rate, kernel_size=1, stride=1, |
| 35 | bias=False)) |
| 36 | self.norm2: nn.BatchNorm2d |
| 37 | self.add_module('norm2', nn.BatchNorm2d(bn_size * growth_rate)) |
| 38 | self.relu2: nn.ReLU |
| 39 | self.add_module('relu2', nn.ReLU(inplace=True)) |
| 40 | self.conv2: nn.Conv2d |
| 41 | self.add_module('conv2', nn.Conv2d(bn_size * growth_rate, growth_rate, |
| 42 | kernel_size=3, stride=1, padding=1, |
| 43 | bias=False)) |
| 44 | self.drop_rate = float(drop_rate) |
| 45 | self.memory_efficient = memory_efficient |
| 46 | |
| 47 | def bn_function(self, inputs: List[Tensor]) -> Tensor: |
| 48 | concated_features = torch.cat(inputs, 1) |