140 lines
4.1 KiB
Python
140 lines
4.1 KiB
Python
import torch
|
|
import numpy as np
|
|
|
|
|
|
def move_to(obj, device):
|
|
if isinstance(obj, torch.Tensor):
|
|
return obj.to(device)
|
|
if isinstance(obj, dict):
|
|
return {k: move_to(v, device) for k, v in obj.items()}
|
|
if isinstance(obj, (list, tuple)):
|
|
return [move_to(x, device) for x in obj]
|
|
return obj # otherwise do nothing
|
|
|
|
|
|
def detach_all(obj):
|
|
if isinstance(obj, torch.Tensor):
|
|
return obj.detach()
|
|
if isinstance(obj, dict):
|
|
return {k: detach_all(v) for k, v in obj.items()}
|
|
if isinstance(obj, (list, tuple)):
|
|
return [detach_all(x) for x in obj]
|
|
return obj # otherwise do nothing
|
|
|
|
|
|
def to_torch(obj):
|
|
if isinstance(obj, np.ndarray):
|
|
return torch.from_numpy(obj).float()
|
|
if isinstance(obj, dict):
|
|
return {k: to_torch(v) for k, v in obj.items()}
|
|
if isinstance(obj, (list, tuple)):
|
|
return [to_torch(x) for x in obj]
|
|
return obj
|
|
|
|
|
|
def to_np(obj):
|
|
if isinstance(obj, torch.Tensor):
|
|
return obj.numpy()
|
|
if isinstance(obj, dict):
|
|
return {k: to_np(v) for k, v in obj.items()}
|
|
if isinstance(obj, (list, tuple)):
|
|
return [to_np(x) for x in obj]
|
|
return obj
|
|
|
|
|
|
def get_device(i=0):
|
|
device = f"cuda:{i}" if torch.cuda.is_available() else "cpu"
|
|
return torch.device(device)
|
|
|
|
|
|
def scatter_intervals(tensor, start, end, T=None):
|
|
"""
|
|
Scatter the tensor contents into intervals from start to end
|
|
output tensor indexed from 0 to end.max()
|
|
:param tensor (B, S, *)
|
|
:param start (B) start indices
|
|
:param end (B) end indices
|
|
:param T (int, optional) max length
|
|
returns (B, T, *) scattered tensor
|
|
"""
|
|
assert isinstance(tensor, torch.Tensor) and tensor.ndim >= 2
|
|
if T is not None:
|
|
end = torch.where(end < 0, T + 1 + end, end)
|
|
assert torch.all(end <= T)
|
|
else:
|
|
T = end.max()
|
|
|
|
B, S, *dims = tensor.shape
|
|
start, end = start.long(), end.long()
|
|
# get idcs that go past the last time step so we don't have repeat indices in scatter
|
|
idcs = time_segment_idcs(start, end, clip=False) # (B, S)
|
|
# mask out the extra padding
|
|
mask = idcs >= end[:, None]
|
|
tensor[mask] = 0
|
|
|
|
idcs = idcs.reshape(B, S, *(1,) * len(dims)).repeat(1, 1, *dims)
|
|
output = torch.zeros(
|
|
B, idcs.max() + 1, *dims, device=tensor.device, dtype=tensor.dtype
|
|
)
|
|
output.scatter_(1, idcs, tensor)
|
|
# slice out the extra segments
|
|
return output[:, :T]
|
|
|
|
|
|
def get_scatter_mask(start, end, T):
|
|
"""
|
|
get the mask of selected intervals
|
|
"""
|
|
B = start.shape[0]
|
|
start, end = start.long(), end.long()
|
|
assert torch.all(end <= T)
|
|
idcs = time_segment_idcs(start, end, clip=True)
|
|
mask = torch.zeros(B, T, device=start.device, dtype=torch.bool)
|
|
mask.scatter_(1, idcs, 1)
|
|
return mask
|
|
|
|
|
|
def select_intervals(series, start, end):
|
|
"""
|
|
Select slices of a tensor from start to end
|
|
will pad uneven sequences to all the max segment length
|
|
:param series (B, T, *)
|
|
:param start (B)
|
|
:param end (B)
|
|
returns (B, S, *) selected segments, S = max(end - start)
|
|
"""
|
|
B, T, *dims = series.shape
|
|
assert torch.all(end <= T)
|
|
sel = time_segment_idcs(start, end, clip=True)
|
|
S = sel.shape[1]
|
|
sel = sel.reshape(B, S, *(1,) * len(dims)).repeat(1, 1, *dims)
|
|
return torch.gather(series, 1, sel)
|
|
|
|
|
|
def get_select_mask(start, end):
|
|
"""
|
|
get the mask of unpadded elementes for the selected time segments
|
|
e.g. sel[mask] are the unpadded elements
|
|
:param start (B)
|
|
:param end (B)
|
|
"""
|
|
idcs = time_segment_idcs(start, end, clip=False)
|
|
return idcs < end[:, None] # (B, S)
|
|
|
|
|
|
def time_segment_idcs(start, end, clip=True):
|
|
"""
|
|
:param start (B)
|
|
:param end (B)
|
|
returns (B, S) long tensor of indices, where S = max(end - start)
|
|
"""
|
|
start, end = start.long(), end.long()
|
|
S = (end - start).max()
|
|
seg = torch.arange(S, dtype=torch.int64, device=start.device)
|
|
idcs = start[:, None] + seg[None, :] # (B, S)
|
|
if clip:
|
|
# clip at the lengths of each track
|
|
imax = end[:, None] - 1
|
|
idcs = idcs.clamp(max=imax)
|
|
return idcs
|