| 7 | add_aio_decoder_specific, add_aio_backbone_specific, add_aio_neck_specific) |
| 8 | |
| 9 | class model_entry(nn.Module): |
| 10 | def __init__(self, backbone_module, neck_module, decoder_module): |
| 11 | super(model_entry, self).__init__() |
| 12 | self.backbone_module = backbone_module |
| 13 | self.neck_module = neck_module |
| 14 | self.decoder_module = decoder_module |
| 15 | add_task_specific(self, False) |
| 16 | add_backbone_specific(self.backbone_module, True) |
| 17 | add_neck_specific(self.neck_module, True) |
| 18 | add_decoder_specific(self.decoder_module, True) |
| 19 | if hasattr(self.decoder_module, 'loss'): |
| 20 | if hasattr(self.decoder_module.loss, 'classifier'): |
| 21 | add_task_specific(self.decoder_module.loss, True) |
| 22 | |
| 23 | def forward(self, input_var, current_step): |
| 24 | x = self.backbone_module(input_var) # {'image': img_mask, 'label': target_mask, 'filename': img_name, 'backbone_output':xxx} |
| 25 | x = self.neck_module(x) |
| 26 | decoder_feature = self.decoder_module(x) |
| 27 | return decoder_feature |
| 28 | |
| 29 | class aio_entry(nn.Module): |
| 30 | def __init__(self, backbone_module, neck_module, decoder_module): |
no outgoing calls