(self, blocks_args=None, global_params=None)
| 112 | """ |
| 113 | |
| 114 | def __init__(self, blocks_args=None, global_params=None): |
| 115 | super().__init__() |
| 116 | assert isinstance(blocks_args, list), 'blocks_args should be a list' |
| 117 | assert len(blocks_args) > 0, 'block args must be greater than 0' |
| 118 | self._global_params = global_params |
| 119 | self._blocks_args = blocks_args |
| 120 | |
| 121 | # Get static or dynamic convolution depending on image size |
| 122 | Conv2d = get_same_padding_conv2d(image_size=global_params.image_size) |
| 123 | |
| 124 | # Batch norm parameters |
| 125 | bn_mom = 1 - self._global_params.batch_norm_momentum |
| 126 | bn_eps = self._global_params.batch_norm_epsilon |
| 127 | |
| 128 | # Stem |
| 129 | in_channels = 3 # rgb |
| 130 | out_channels = round_filters(32, self._global_params) # number of output channels |
| 131 | self._conv_stem = Conv2d(in_channels, out_channels, kernel_size=3, stride=2, bias=False) |
| 132 | self._bn0 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps) |
| 133 | |
| 134 | # Build blocks |
| 135 | self._blocks = nn.ModuleList([]) |
| 136 | for block_args in self._blocks_args: |
| 137 | |
| 138 | # Update block input and output filters based on depth multiplier. |
| 139 | block_args = block_args._replace( |
| 140 | input_filters=round_filters(block_args.input_filters, self._global_params), |
| 141 | output_filters=round_filters(block_args.output_filters, self._global_params), |
| 142 | num_repeat=round_repeats(block_args.num_repeat, self._global_params) |
| 143 | ) |
| 144 | |
| 145 | # The first block needs to take care of stride and filter size increase. |
| 146 | self._blocks.append(MBConvBlock(block_args, self._global_params)) |
| 147 | if block_args.num_repeat > 1: |
| 148 | block_args = block_args._replace(input_filters=block_args.output_filters, stride=1) |
| 149 | for _ in range(block_args.num_repeat - 1): |
| 150 | self._blocks.append(MBConvBlock(block_args, self._global_params)) |
| 151 | |
| 152 | # Head |
| 153 | in_channels = block_args.output_filters # output of final block |
| 154 | out_channels = round_filters(1280, self._global_params) |
| 155 | self._conv_head = Conv2d(in_channels, out_channels, kernel_size=1, bias=False) |
| 156 | self._bn1 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps) |
| 157 | |
| 158 | # Final linear layer |
| 159 | self._avg_pooling = nn.AdaptiveAvgPool2d(1) |
| 160 | self._dropout = nn.Dropout(self._global_params.dropout_rate) |
| 161 | self._fc = nn.Linear(out_channels, self._global_params.num_classes) |
| 162 | self._swish = MemoryEfficientSwish() |
| 163 | |
| 164 | def set_swish(self, memory_efficient=True): |
| 165 | """Sets swish function as memory efficient (for training) or standard (for export)""" |
nothing calls this directly
no test coverage detected