Hierarchical attention, primarily used for getting best inference results. We use attention at multiple scales, giving priority to the lower resolutions. For example, if we have 4 scales {0.5, 1.0, 1.5, 2.0}, then evaluation is done as follows:
(self, inputs, scales)
| 183 | 'logit_attn': attn} |
| 184 | |
| 185 | def nscale_forward(self, inputs, scales): |
| 186 | """ |
| 187 | Hierarchical attention, primarily used for getting best inference |
| 188 | results. |
| 189 | |
| 190 | We use attention at multiple scales, giving priority to the lower |
| 191 | resolutions. For example, if we have 4 scales {0.5, 1.0, 1.5, 2.0}, |
| 192 | then evaluation is done as follows: |
| 193 | |
| 194 | p_joint = attn_1.5 * p_1.5 + (1 - attn_1.5) * down(p_2.0) |
| 195 | p_joint = attn_1.0 * p_1.0 + (1 - attn_1.0) * down(p_joint) |
| 196 | p_joint = up(attn_0.5 * p_0.5) * (1 - up(attn_0.5)) * p_joint |
| 197 | |
| 198 | The target scale is always 1.0, and 1.0 is expected to be part of the |
| 199 | list of scales. When predictions are done at greater than 1.0 scale, |
| 200 | the predictions are downsampled before combining with the next lower |
| 201 | scale. |
| 202 | |
| 203 | Inputs: |
| 204 | scales - a list of scales to evaluate |
| 205 | inputs - dict containing 'images', the input, and 'gts', the ground |
| 206 | truth mask |
| 207 | |
| 208 | Output: |
| 209 | If training, return loss, else return prediction + attention |
| 210 | """ |
| 211 | x_1x = inputs['images'] |
| 212 | |
| 213 | assert 1.0 in scales, 'expected 1.0 to be the target scale' |
| 214 | # Lower resolution provides attention for higher rez predictions, |
| 215 | # so we evaluate in order: high to low |
| 216 | scales = sorted(scales, reverse=True) |
| 217 | |
| 218 | pred = None |
| 219 | aux = None |
| 220 | output_dict = {} |
| 221 | |
| 222 | for s in scales: |
| 223 | x = ResizeX(x_1x, s) |
| 224 | outs = self._fwd(x) |
| 225 | cls_out = outs['cls_out'] |
| 226 | attn_out = outs['logit_attn'] |
| 227 | aux_out = outs['aux_out'] |
| 228 | |
| 229 | output_dict[fmt_scale('pred', s)] = cls_out |
| 230 | if s != 2.0: |
| 231 | output_dict[fmt_scale('attn', s)] = attn_out |
| 232 | |
| 233 | if pred is None: |
| 234 | pred = cls_out |
| 235 | aux = aux_out |
| 236 | elif s >= 1.0: |
| 237 | # downscale previous |
| 238 | pred = scale_as(pred, cls_out) |
| 239 | pred = attn_out * cls_out + (1 - attn_out) * pred |
| 240 | aux = scale_as(aux, cls_out) |
| 241 | aux = attn_out * aux_out + (1 - attn_out) * aux |
| 242 | else: |