| 9 | from torch import Tensor |
| 10 | |
| 11 | class _DenseLayer(nn.Module): |
| 12 | def __init__(self, input_c, growth_rate, bn_size, drop_rate, memory_efficient=False): |
| 13 | super(_DenseLayer, self).__init__() |
| 14 | self.add_module("normal", nn.BatchNorm2d(input_c)) |
| 15 | self.add_module("relu1", nn.ReLU(inplace=True)) |
| 16 | self.add_module("conv1", nn.Conv2d(in_channels=input_c, out_channels=bn_size, kernel_size=1, stride=1, bias=False)) |
| 17 | self.add_module("norm2", nn.BatchNorm2d(bn_size * growth_rate)) |
| 18 | self.add_module("conv2", nn.Conv2d(bn_size*growth_rate, growth_rate, kernel_size=3, stride=1, padding=1, bias=False)) |
| 19 | self.dro_rate = drop_rate |
| 20 | self.memory_efficient = memory_efficient |
| 21 | |
| 22 | def bn_function(self, inputs): |
| 23 | concat_features = torch.cat(inputs, 1) |
| 24 | bottleneck_output = self.conv1(self.relu1(self.norm1(concat_features))) |
| 25 | return bottleneck_output |
| 26 | |
| 27 | @staticmethod |
| 28 | def any_requires_grad(self, inputs): |
| 29 | for tensor in inputs: |
| 30 | if tensor.requires_grad: |
| 31 | return True |
| 32 | |
| 33 | return False |
| 34 | |
| 35 | @torch.jit.unused |
| 36 | def call_checkpoint_bottleneck(self, inputs): |
| 37 | def closure(*inp): |
| 38 | return self.bn_function(inp) |
| 39 | return cp.ckeckpoints(closure, *input) |
| 40 | |
| 41 | |
| 42 | def forward(self, inputs): |
| 43 | if isinstance(inputs, Tensor): |
| 44 | prev_features = [inputs] |
| 45 | else: |
| 46 | prev_features = inputs |
| 47 | |
| 48 | if self.memory_efficient and self.any_requires_grad(prev_features): |
| 49 | if torch.jit.is_scripting(): |
| 50 | raise Exception("memory efficient not supported in JIT") |
| 51 | |
| 52 | bottleneck_ouput = self.call_checkpoint_bottleneck(prev_features) |
| 53 | else: |
| 54 | bottleneck_ouput = self.bn_function(prev_features) |
| 55 | |
| 56 | new_features = self.conv2(self.relu2(self.norm2(bottleneck_ouput))) |
| 57 | if self.drop_rate > 0: |
| 58 | new_features = F.dropout(new_features, p=self.drop_rate, training = self.training) |
| 59 | return new_features |
| 60 | |
| 61 | class _DenseBlock(nn.ModuleDict): |
| 62 | _version=2 |