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)
| 112 | return {'pred': pred, 'attn_10x': attn} |
| 113 | |
| 114 | def nscale_forward(self, inputs, scales): |
| 115 | """ |
| 116 | Hierarchical attention, primarily used for getting best inference |
| 117 | results. |
| 118 | |
| 119 | We use attention at multiple scales, giving priority to the lower |
| 120 | resolutions. For example, if we have 4 scales {0.5, 1.0, 1.5, 2.0}, |
| 121 | then evaluation is done as follows: |
| 122 | |
| 123 | p_joint = attn_1.5 * p_1.5 + (1 - attn_1.5) * down(p_2.0) |
| 124 | p_joint = attn_1.0 * p_1.0 + (1 - attn_1.0) * down(p_joint) |
| 125 | p_joint = up(attn_0.5 * p_0.5) * (1 - up(attn_0.5)) * p_joint |
| 126 | |
| 127 | The target scale is always 1.0, and 1.0 is expected to be part of the |
| 128 | list of scales. When predictions are done at greater than 1.0 scale, |
| 129 | the predictions are downsampled before combining with the next lower |
| 130 | scale. |
| 131 | |
| 132 | Inputs: |
| 133 | scales - a list of scales to evaluate |
| 134 | inputs - dict containing 'images', the input, and 'gts', the ground |
| 135 | truth mask |
| 136 | |
| 137 | Output: |
| 138 | If training, return loss, else return prediction + attention |
| 139 | """ |
| 140 | x_1x = inputs['images'] |
| 141 | |
| 142 | assert 1.0 in scales, 'expected 1.0 to be the target scale' |
| 143 | # Lower resolution provides attention for higher rez predictions, |
| 144 | # so we evaluate in order: high to low |
| 145 | scales = sorted(scales, reverse=True) |
| 146 | |
| 147 | pred = None |
| 148 | output_dict = {} |
| 149 | |
| 150 | for s in scales: |
| 151 | x = ResizeX(x_1x, s) |
| 152 | bs = x.shape[0] |
| 153 | scale_float = torch.Tensor(bs).fill_(s) |
| 154 | p, attn, _aspp_attn, _aspp = self._fwd(x, scale_float=scale_float) |
| 155 | |
| 156 | output_dict[fmt_scale('pred', s)] = p |
| 157 | if s != 2.0: |
| 158 | output_dict[fmt_scale('attn', s)] = attn |
| 159 | |
| 160 | if pred is None: |
| 161 | pred = p |
| 162 | elif s >= 1.0: |
| 163 | # downscale previous |
| 164 | pred = scale_as(pred, p) |
| 165 | pred = attn * p + (1 - attn) * pred |
| 166 | else: |
| 167 | # upscale current |
| 168 | p = attn * p |
| 169 | p = scale_as(p, pred) |
| 170 | attn = scale_as(attn, pred) |
| 171 | pred = p + (1 - attn) * pred |