-
Notifications
You must be signed in to change notification settings - Fork 51
/
loss.py
executable file
·37 lines (26 loc) · 1022 Bytes
/
loss.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
import torch
def flow_loss_func(flow_preds, flow_gt, valid,
gamma=0.9,
max_flow=400,
**kwargs,
):
n_predictions = len(flow_preds)
flow_loss = 0.0
# exlude invalid pixels and extremely large diplacements
mag = torch.sum(flow_gt ** 2, dim=1).sqrt() # [B, H, W]
valid = (valid >= 0.5) & (mag < max_flow)
for i in range(n_predictions):
i_weight = gamma ** (n_predictions - i - 1)
i_loss = (flow_preds[i] - flow_gt).abs()
flow_loss += i_weight * (valid[:, None] * i_loss).mean()
epe = torch.sum((flow_preds[-1] - flow_gt) ** 2, dim=1).sqrt()
if valid.max() < 0.5:
pass
epe = epe.view(-1)[valid.view(-1)]
metrics = {
'epe': epe.mean().item(),
'1px': (epe > 1).float().mean().item(),
'3px': (epe > 3).float().mean().item(),
'5px': (epe > 5).float().mean().item(),
}
return flow_loss, metrics