| 173 | |
| 174 | @torch.no_grad() |
| 175 | def compute_coordinates(self, x): |
| 176 | h, w = x.size(2), x.size(3) |
| 177 | y_loc = -1.0 + 2.0 * torch.arange(h, device=x.device) / (h - 1) |
| 178 | x_loc = -1.0 + 2.0 * torch.arange(w, device=x.device) / (w - 1) |
| 179 | y_loc, x_loc = torch.meshgrid(y_loc, x_loc) |
| 180 | y_loc = y_loc.expand([x.shape[0], 1, -1, -1]) |
| 181 | x_loc = x_loc.expand([x.shape[0], 1, -1, -1]) |
| 182 | locations = torch.cat([x_loc, y_loc], 1) |
| 183 | return locations.to(x) |
| 184 | |
| 185 | def forward(self, features, is_training=True, **kwargs): |
| 186 | output = {} |