-
Notifications
You must be signed in to change notification settings - Fork 10
/
utils.py
77 lines (70 loc) · 2.26 KB
/
utils.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
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
import torch
import random
import numpy as np
def ADE_FDE(y_, y, batch_first=False):
# average displacement error
# final displacement error
# y_, y: S x L x N x 2
if torch.is_tensor(y):
err = (y_ - y).norm(dim=-1)
else:
err = np.linalg.norm(np.subtract(y_, y), axis=-1)
if len(err.shape) == 1:
fde = err[-1]
ade = err.mean()
elif batch_first:
fde = err[..., -1]
ade = err.mean(-1)
else:
fde = err[..., -1, :]
ade = err.mean(-2)
return ade, fde
def kmeans(k, data, iters=None):
centroids = data.copy()
np.random.shuffle(centroids)
centroids = centroids[:k]
if iters is None: iters = 100000
for _ in range(iters):
# while True:
distances = np.sqrt(((data - centroids[:, np.newaxis])**2).sum(axis=2))
closest = np.argmin(distances, axis=0)
centroids_ = []
for k in range(len(centroids)):
cand = data[closest==k]
if len(cand) > 0:
centroids_.append(cand.mean(axis=0))
else:
centroids_.append(data[np.random.randint(len(data))])
centroids_ = np.array(centroids_)
if np.linalg.norm(centroids_ - centroids) < 0.0001:
break
centroids = centroids_
return centroids
def FPC(y, n_samples):
# y: S x L x 2
goal = y[...,-1,:2]
goal_ = kmeans(n_samples, goal)
dist = np.linalg.norm(goal_[:,np.newaxis,:2] - goal[np.newaxis,:,:2], axis=-1)
chosen = np.argmin(dist, axis=1)
return chosen
def seed(seed: int):
rand = seed is None
if seed is None:
seed = int.from_bytes(os.urandom(4), byteorder="big")
np.random.seed(seed)
random.seed(seed)
torch.manual_seed(seed)
torch.backends.cudnn.deterministic = not rand
torch.backends.cudnn.benchmark = rand
def get_rng_state(device):
return (
torch.get_rng_state(),
torch.cuda.get_rng_state(device) if torch.cuda.is_available and "cuda" in str(device) else None,
np.random.get_state(),
random.getstate(),
)
def set_rng_state(state, device):
torch.set_rng_state(state[0])
if state[1] is not None: torch.cuda.set_rng_state(state[1], device)
np.random.set_state(state[2])
random.setstate(state[3])