OCR net
| 92 | |
| 93 | |
| 94 | class OCRNet(nn.Module): |
| 95 | """ |
| 96 | OCR net |
| 97 | """ |
| 98 | def __init__(self, num_classes, trunk='hrnetv2', criterion=None): |
| 99 | super(OCRNet, self).__init__() |
| 100 | self.criterion = criterion |
| 101 | self.backbone, _, _, high_level_ch = get_trunk(trunk) |
| 102 | self.ocr = OCR_block(high_level_ch) |
| 103 | |
| 104 | def forward(self, inputs): |
| 105 | assert 'images' in inputs |
| 106 | x = inputs['images'] |
| 107 | |
| 108 | _, _, high_level_features = self.backbone(x) |
| 109 | cls_out, aux_out, _ = self.ocr(high_level_features) |
| 110 | aux_out = scale_as(aux_out, x) |
| 111 | cls_out = scale_as(cls_out, x) |
| 112 | |
| 113 | if self.training: |
| 114 | gts = inputs['gts'] |
| 115 | aux_loss = self.criterion(aux_out, gts, |
| 116 | do_rmi=cfg.LOSS.OCR_AUX_RMI) |
| 117 | main_loss = self.criterion(cls_out, gts) |
| 118 | loss = cfg.LOSS.OCR_ALPHA * aux_loss + main_loss |
| 119 | return loss |
| 120 | else: |
| 121 | output_dict = {'pred': cls_out} |
| 122 | return output_dict |
| 123 | |
| 124 | |
| 125 | class OCRNetASPP(nn.Module): |