| 95 | |
| 96 | |
| 97 | class _DenseBlock(nn.ModuleDict): |
| 98 | _version = 2 |
| 99 | |
| 100 | def __init__( |
| 101 | self, |
| 102 | num_layers: int, |
| 103 | num_input_features: int, |
| 104 | bn_size: int, |
| 105 | growth_rate: int, |
| 106 | drop_rate: float, |
| 107 | memory_efficient: bool = False |
| 108 | ) -> None: |
| 109 | super(_DenseBlock, self).__init__() |
| 110 | for i in range(num_layers): |
| 111 | layer = _DenseLayer( |
| 112 | num_input_features + i * growth_rate, |
| 113 | growth_rate=growth_rate, |
| 114 | bn_size=bn_size, |
| 115 | drop_rate=drop_rate, |
| 116 | memory_efficient=memory_efficient, |
| 117 | ) |
| 118 | self.add_module('denselayer%d' % (i + 1), layer) |
| 119 | |
| 120 | def forward(self, init_features: Tensor) -> Tensor: |
| 121 | features = [init_features] |
| 122 | for name, layer in self.items(): |
| 123 | new_features = layer(features) |
| 124 | features.append(new_features) |
| 125 | return torch.cat(features, 1) |
| 126 | |
| 127 | |
| 128 | class _Transition(nn.Sequential): |