| 16 | |
| 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) |
| 49 | bottleneck_output = self.conv1(self.relu1(self.norm1(concated_features))) # noqa: T484 |
| 50 | return bottleneck_output |
| 51 | |
| 52 | # todo: rewrite when torchscript supports any |
| 53 | def any_requires_grad(self, input: List[Tensor]) -> bool: |
| 54 | for tensor in input: |
| 55 | if tensor.requires_grad: |
| 56 | return True |
| 57 | return False |
| 58 | |
| 59 | @torch.jit.unused # noqa: T484 |
| 60 | def call_checkpoint_bottleneck(self, input: List[Tensor]) -> Tensor: |
| 61 | def closure(*inputs): |
| 62 | return self.bn_function(inputs) |
| 63 | |
| 64 | return cp.checkpoint(closure, *input) |
| 65 | |
| 66 | @torch.jit._overload_method # noqa: F811 |
| 67 | def forward(self, input: List[Tensor]) -> Tensor: |
| 68 | pass |
| 69 | |
| 70 | @torch.jit._overload_method # noqa: F811 |
| 71 | def forward(self, input: Tensor) -> Tensor: |
| 72 | pass |
| 73 | |
| 74 | # torchscript does not yet support *args, so we overload method |
| 75 | # allowing it to take either a List[Tensor] or single Tensor |