| 127 | 'attn_10x': attn} |
| 128 | |
| 129 | def two_scale_forward(self, inputs): |
| 130 | assert 'images' in inputs |
| 131 | |
| 132 | x_1x = inputs['images'] |
| 133 | x_lo = ResizeX(x_1x, cfg.MODEL.MSCALE_LO_SCALE) |
| 134 | |
| 135 | p_lo, feats_lo = self._fwd(x_lo) |
| 136 | p_1x, feats_hi = self._fwd(x_1x) |
| 137 | |
| 138 | feats_hi = scale_as(feats_hi, feats_lo) |
| 139 | cat_feats = torch.cat([feats_lo, feats_hi], 1) |
| 140 | logit_attn = self.scale_attn(cat_feats) |
| 141 | logit_attn = scale_as(logit_attn, p_lo) |
| 142 | |
| 143 | p_lo = logit_attn * p_lo |
| 144 | p_lo = scale_as(p_lo, p_1x) |
| 145 | logit_attn = scale_as(logit_attn, p_1x) |
| 146 | joint_pred = p_lo + (1 - logit_attn) * p_1x |
| 147 | |
| 148 | if self.training: |
| 149 | assert 'gts' in inputs |
| 150 | gts = inputs['gts'] |
| 151 | loss = self.criterion(joint_pred, gts) |
| 152 | return loss |
| 153 | else: |
| 154 | # FIXME: should add multi-scale values for pred and attn |
| 155 | return {'pred': joint_pred, |
| 156 | 'attn_10x': logit_attn} |
| 157 | |
| 158 | def forward(self, inputs): |
| 159 | if cfg.MODEL.N_SCALES and not self.training: |