(self)
| 156 | pin_memory=False, sampler=self.sampler) |
| 157 | |
| 158 | def create_model(self): |
| 159 | config = self.config |
| 160 | ginfo = self.ginfo |
| 161 | |
| 162 | backbone_bn_group_size = config.backbone['kwargs'].get('bn_group_size', 1) |
| 163 | assert backbone_bn_group_size == 1, 'other bn group size not support!' |
| 164 | backbone_bn_group_comm = self.ginfo.backbone_share_group |
| 165 | # if backbone_bn_group_size == 1: |
| 166 | # backbone_bn_group_comm = None |
| 167 | # else: |
| 168 | # assert self.C.world_size % backbone_bn_group_size == 0 |
| 169 | # backbone_bn_group_comm = simple_group_split(self.C.world_size, self.C.rank, self.C.world_size // backbone_bn_group_size) |
| 170 | |
| 171 | ## build backbone |
| 172 | config.backbone['kwargs']['bn_group'] = backbone_bn_group_comm |
| 173 | backbone_module = backbones.backbone_entry(config.backbone) |
| 174 | count_parameters_num(backbone_module) |
| 175 | |
| 176 | ## build necks |
| 177 | neck_bn_group_size = config.backbone['kwargs'].get('bn_group_size', 1) |
| 178 | assert neck_bn_group_size == 1, 'other bn group size not support!' |
| 179 | neck_bn_group_comm = self.ginfo.neck_share_group |
| 180 | |
| 181 | # neck_bn_group_size = config.neck['kwargs'].get('bn_group_size', 1) |
| 182 | # if neck_bn_group_size == 1: |
| 183 | # neck_bn_group_comm = None |
| 184 | # else: |
| 185 | # assert self.C.world_size % neck_bn_group_size == 0 |
| 186 | # neck_bn_group_comm = simple_group_split(self.C.world_size, self.C.rank, self.C.world_size // neck_bn_group_size) |
| 187 | |
| 188 | neck_args = config.neck['kwargs'] |
| 189 | neck_args['backbone'] = backbone_module |
| 190 | neck_args['bn_group'] = neck_bn_group_comm |
| 191 | neck_module = necks.neck_entry(config.neck) |
| 192 | |
| 193 | ## add decoder |
| 194 | decoder_bn_group_size = config.backbone['kwargs'].get('bn_group_size', 1) |
| 195 | assert decoder_bn_group_size == 1, 'other bn group size not support!' |
| 196 | decoder_bn_group_comm = self.ginfo.decoder_share_group |
| 197 | |
| 198 | # decoder_bn_group_size = config.neck['kwargs'].get('bn_group_size', 1) |
| 199 | # if decoder_bn_group_size == 1: |
| 200 | # decoder_bn_group_comm = None |
| 201 | # else: |
| 202 | # assert self.C.world_size % decoder_bn_group_size == 0 |
| 203 | # decoder_bn_group_comm = simple_group_split(self.C.world_size, self.C.rank, self.C.world_size // decoder_bn_group_size) |
| 204 | |
| 205 | decoder_args = config.decoder['kwargs'] |
| 206 | decoder_args['backbone'] = backbone_module |
| 207 | decoder_args['neck'] = neck_module |
| 208 | decoder_args['bn_group'] = decoder_bn_group_comm |
| 209 | decoder_module = decoders.decoder_entry(config.decoder) |
| 210 | |
| 211 | # build |
| 212 | model = model_entry(backbone_module, neck_module, decoder_module) |
| 213 | |
| 214 | if self.C.rank == 0: |
| 215 | print(model) |
no test coverage detected