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)
| 53 | pass |
| 54 | |
| 55 | def nscale_forward(self, inputs, scales): |
| 56 | """ |
| 57 | Hierarchical attention, primarily used for getting best inference |
| 58 | results. |
| 59 | |
| 60 | We use attention at multiple scales, giving priority to the lower |
| 61 | resolutions. For example, if we have 4 scales {0.5, 1.0, 1.5, 2.0}, |
| 62 | then evaluation is done as follows: |
| 63 | |
| 64 | p_joint = attn_1.5 * p_1.5 + (1 - attn_1.5) * down(p_2.0) |
| 65 | p_joint = attn_1.0 * p_1.0 + (1 - attn_1.0) * down(p_joint) |
| 66 | p_joint = up(attn_0.5 * p_0.5) * (1 - up(attn_0.5)) * p_joint |
| 67 | |
| 68 | The target scale is always 1.0, and 1.0 is expected to be part of the |
| 69 | list of scales. When predictions are done at greater than 1.0 scale, |
| 70 | the predictions are downsampled before combining with the next lower |
| 71 | scale. |
| 72 | |
| 73 | Inputs: |
| 74 | scales - a list of scales to evaluate |
| 75 | inputs - dict containing 'images', the input, and 'gts', the ground |
| 76 | truth mask |
| 77 | |
| 78 | Output: |
| 79 | If training, return loss, else return prediction + attention |
| 80 | """ |
| 81 | x_1x = inputs['images'] |
| 82 | |
| 83 | assert 1.0 in scales, 'expected 1.0 to be the target scale' |
| 84 | # Lower resolution provides attention for higher rez predictions, |
| 85 | # so we evaluate in order: high to low |
| 86 | scales = sorted(scales, reverse=True) |
| 87 | pred = None |
| 88 | last_feats = None |
| 89 | |
| 90 | for idx, s in enumerate(scales): |
| 91 | x = ResizeX(x_1x, s) |
| 92 | p, feats = self._fwd(x) |
| 93 | |
| 94 | # Generate attention prediction |
| 95 | if idx > 0: |
| 96 | assert last_feats is not None |
| 97 | # downscale feats |
| 98 | last_feats = scale_as(last_feats, feats) |
| 99 | cat_feats = torch.cat([feats, last_feats], 1) |
| 100 | attn = self.scale_attn(cat_feats) |
| 101 | attn = scale_as(attn, p) |
| 102 | |
| 103 | if pred is None: |
| 104 | # This is the top scale prediction |
| 105 | pred = p |
| 106 | elif s >= 1.0: |
| 107 | # downscale previous |
| 108 | pred = scale_as(pred, p) |
| 109 | pred = attn * p + (1 - attn) * pred |
| 110 | else: |
| 111 | # upscale current |
| 112 | p = attn * p |