| 115 | |
| 116 | |
| 117 | class _DenseBlock(nn.ModuleDict): |
| 118 | _version = 2 |
| 119 | |
| 120 | def __init__(self, num_layers, num_input_features, bn_size, growth_rate, drop_rate, memory_efficient=False): |
| 121 | super(_DenseBlock, self).__init__() |
| 122 | for i in range(num_layers): |
| 123 | layer = _DenseLayer( |
| 124 | num_input_features + i * growth_rate, |
| 125 | growth_rate=growth_rate, |
| 126 | bn_size=bn_size, |
| 127 | drop_rate=drop_rate, |
| 128 | memory_efficient=memory_efficient, |
| 129 | ) |
| 130 | self.add_module('denselayer%d' % (i + 1), layer) |
| 131 | |
| 132 | def forward(self, init_features): |
| 133 | features = [init_features] |
| 134 | for name, layer in self.items(): |
| 135 | new_features = layer(features) |
| 136 | features.append(new_features) |
| 137 | return torch.cat(features, 1) |
| 138 | |
| 139 | |
| 140 | class _Transition(nn.Sequential): |