Do we supervised both aux outputs, lo and high scale? Should attention be used to combine the aux output? Normally we only supervise the combined 1x output If we use attention to combine the aux outputs, then we can use normal weighting for aux vs. cls outpu
(self, inputs)
| 262 | return output_dict |
| 263 | |
| 264 | def two_scale_forward(self, inputs): |
| 265 | """ |
| 266 | Do we supervised both aux outputs, lo and high scale? |
| 267 | Should attention be used to combine the aux output? |
| 268 | Normally we only supervise the combined 1x output |
| 269 | |
| 270 | If we use attention to combine the aux outputs, then |
| 271 | we can use normal weighting for aux vs. cls outputs |
| 272 | """ |
| 273 | assert 'images' in inputs |
| 274 | x_1x = inputs['images'] |
| 275 | |
| 276 | x_lo = ResizeX(x_1x, cfg.MODEL.MSCALE_LO_SCALE) |
| 277 | lo_outs = self._fwd(x_lo) |
| 278 | pred_05x = lo_outs['cls_out'] |
| 279 | p_lo = pred_05x |
| 280 | aux_lo = lo_outs['aux_out'] |
| 281 | logit_attn = lo_outs['logit_attn'] |
| 282 | attn_05x = logit_attn |
| 283 | |
| 284 | hi_outs = self._fwd(x_1x) |
| 285 | pred_10x = hi_outs['cls_out'] |
| 286 | p_1x = pred_10x |
| 287 | aux_1x = hi_outs['aux_out'] |
| 288 | |
| 289 | p_lo = logit_attn * p_lo |
| 290 | aux_lo = logit_attn * aux_lo |
| 291 | p_lo = scale_as(p_lo, p_1x) |
| 292 | aux_lo = scale_as(aux_lo, p_1x) |
| 293 | |
| 294 | logit_attn = scale_as(logit_attn, p_1x) |
| 295 | |
| 296 | # combine lo and hi predictions with attention |
| 297 | joint_pred = p_lo + (1 - logit_attn) * p_1x |
| 298 | joint_aux = aux_lo + (1 - logit_attn) * aux_1x |
| 299 | |
| 300 | if self.training: |
| 301 | gts = inputs['gts'] |
| 302 | do_rmi = cfg.LOSS.OCR_AUX_RMI |
| 303 | aux_loss = self.criterion(joint_aux, gts, do_rmi=do_rmi) |
| 304 | |
| 305 | # Optionally turn off RMI loss for first epoch to try to work |
| 306 | # around cholesky errors of singular matrix |
| 307 | do_rmi_main = True # cfg.EPOCH > 0 |
| 308 | main_loss = self.criterion(joint_pred, gts, do_rmi=do_rmi_main) |
| 309 | loss = cfg.LOSS.OCR_ALPHA * aux_loss + main_loss |
| 310 | |
| 311 | # Optionally, apply supervision to the multi-scale predictions |
| 312 | # directly. Turn off RMI to keep things lightweight |
| 313 | if cfg.LOSS.SUPERVISED_MSCALE_WT: |
| 314 | scaled_pred_05x = scale_as(pred_05x, p_1x) |
| 315 | loss_lo = self.criterion(scaled_pred_05x, gts, do_rmi=False) |
| 316 | loss_hi = self.criterion(pred_10x, gts, do_rmi=False) |
| 317 | loss += cfg.LOSS.SUPERVISED_MSCALE_WT * loss_lo |
| 318 | loss += cfg.LOSS.SUPERVISED_MSCALE_WT * loss_hi |
| 319 | return loss |
| 320 | else: |
| 321 | output_dict = { |