| 44 | |
| 45 | |
| 46 | class _DenseLayer(nn.Module): |
| 47 | def __init__(self, num_input_features, growth_rate, bn_size, drop_rate, memory_efficient=False): |
| 48 | super(_DenseLayer, self).__init__() |
| 49 | self.add_module('norm1', nn.BatchNorm2d(num_input_features)), |
| 50 | self.add_module('relu1', nn.ReLU(inplace=False)), |
| 51 | self.add_module('conv1', nn.Conv2d(num_input_features, bn_size * |
| 52 | growth_rate, kernel_size=1, stride=1, |
| 53 | bias=False)), |
| 54 | self.add_module('norm2', nn.BatchNorm2d(bn_size * growth_rate)), |
| 55 | self.add_module('relu2', nn.ReLU(inplace=False)), |
| 56 | self.add_module('conv2', nn.Conv2d(bn_size * growth_rate, growth_rate, |
| 57 | kernel_size=3, stride=1, padding=1, |
| 58 | bias=False)), |
| 59 | self.drop_rate = float(drop_rate) |
| 60 | self.memory_efficient = memory_efficient |
| 61 | |
| 62 | def bn_function(self, inputs): |
| 63 | # type: (List[Tensor]) -> Tensor |
| 64 | concated_features = torch.cat(inputs, 1) |
| 65 | bottleneck_output = self.conv1(self.relu1(self.norm1(concated_features))) # noqa: T484 |
| 66 | return bottleneck_output |
| 67 | |
| 68 | # todo: rewrite when torchscript supports any |
| 69 | def any_requires_grad(self, input): |
| 70 | # type: (List[Tensor]) -> bool |
| 71 | for tensor in input: |
| 72 | if tensor.requires_grad: |
| 73 | return True |
| 74 | return False |
| 75 | |
| 76 | @torch.jit.unused # noqa: T484 |
| 77 | def call_checkpoint_bottleneck(self, input): |
| 78 | # type: (List[Tensor]) -> Tensor |
| 79 | def closure(*inputs): |
| 80 | return self.bn_function(inputs) |
| 81 | |
| 82 | return cp.checkpoint(closure, *input) |
| 83 | |
| 84 | @torch.jit._overload_method # noqa: F811 |
| 85 | def forward(self, input): |
| 86 | # type: (List[Tensor]) -> (Tensor) |
| 87 | pass |
| 88 | |
| 89 | @torch.jit._overload_method # noqa: F811 |
| 90 | def forward(self, input): |
| 91 | # type: (Tensor) -> (Tensor) |
| 92 | pass |
| 93 | |
| 94 | # torchscript does not yet support *args, so we overload method |
| 95 | # allowing it to take either a List[Tensor] or single Tensor |
| 96 | def forward(self, input): # noqa: F811 |
| 97 | if isinstance(input, Tensor): |
| 98 | prev_features = [input] |
| 99 | else: |
| 100 | prev_features = input |
| 101 | |
| 102 | if self.memory_efficient and self.any_requires_grad(prev_features): |
| 103 | if torch.jit.is_scripting(): |