def augmented_dynamics(t, y_aug):
# Dynamics of the original system augmented with
# the adjoint wrt y, and an integrator wrt t and args.
y = y_aug[1]
adj_y = y_aug[2]
# ignore gradients wrt time and parameters
with torch.enable_grad():
t_ = t.detach()
t = t_.requires_grad_(True)
y = y.detach().requires_grad_(True)
# If using an adaptive solver we don't want to waste time resolving dL/dt unless we need it (which
# doesn't necessarily even exist if there is piecewise structure in time), so turning off gradients
# wrt t here means we won't compute that if we don't need it.
func_eval = func(t if t_requires_grad else t_, y)
# Workaround for PyTorch bug #39784
_t = torch.as_strided(t, (), ())
_y = torch.as_strided(y, (), ())
_params = tuple(torch.as_strided(param, (), ()) for param in adjoint_params)
vjp_t, vjp_y, *vjp_params = torch.autograd.grad(
func_eval, (t, y) + adjoint_params, -adj_y,
allow_unused=True, retain_graph=True
)
# autograd.grad returns None if no gradient, set to zero.
vjp_t = torch.zeros_like(t) if vjp_t is None else vjp_t
vjp_y = torch.zeros_like(y) if vjp_y is None else vjp_y
vjp_params = [torch.zeros_like(param) if vjp_param is None else vjp_param
for param, vjp_param in zip(adjoint_params, vjp_params)]
return (vjp_t, func_eval, vjp_y, *vjp_params)
def complex_view(x, dim=-1, squeeze=True):
r"""Returns a real and imaginary views into the complex tensor, assumed
to be in interleaved layout (double-real, i.e. re-im). The returned tensor
is half the size along the specified dimension and is not, in general,
contiguous in memory.
Arguments
---------
x : torch.tensor
Time series of measurement values
dim : int
Axis along which the periodogram is computed, i.e. ``dim=-1``.
Returns
-------
real : torch.tensor
The view into a real part of the tensor.
imag : torch.tensor
The view into a real part of the tensor.
"""
dim = fix_dim(dim, x.dim())
shape, strides = list(x.size()), list(x.stride())
offset = x.storage_offset()
# compute new shape and strides
strided_size, rem = divmod(shape[dim], 2)
if rem != 0:
warnings.warn(
f"Odd dimension size for the complex data unpacking: "
f"taking the least size that fits.", RuntimeWarning)
# new shape and stride structure
if shape[dim] == 2 and squeeze:
# if the complex dimension is exactly two, then just drop it
shape_view = shape[:dim] + shape[dim + 1:]
strides_view = strides[:dim] + strides[dim + 1:]
else:
# otherwise, half the size and double the stride
size, rem = divmod(shape[dim], 2)
shape_view = shape[:dim] + [size] + shape[dim + 1:]
strides_view = strides[:dim] + [2 * strides[dim]] + strides[dim + 1:]
# differentiable strided view into real and imaginary parts
real = torch.as_strided(x, shape_view, strides_view, offset)
imag = torch.as_strided(x, shape_view, strides_view, offset + strides[dim])
return real, imag
def preprocess_audio_batch(audio, sr, center=True, hop_size=0.1):
if audio.dim() == 2:
audio = torch.mean(audio, axis=1)
if sr != TARGET_SR:
audio = julius.resample_frac(audio, sr, TARGET_SR)
audio_len = audio.size()[0]
frame_len = TARGET_SR
hop_len = int(hop_size * TARGET_SR)
if center:
audio = center_audio(audio, frame_len)
audio = pad_audio(audio, frame_len, hop_len)
n_frames = 1 + int((len(audio) - frame_len) / float(hop_len))
x = torch.as_strided(
audio,
size=(frame_len, n_frames),
stride=(1, hop_len),
)
x = torch.transpose(x, 0, 1)
x = x.unsqueeze(1)
return x
def gen_sequence_batch(batch_size: int,
input_size: int,
min_len: int,
max_len: int,
device: str = 'cpu') -> List[Tensor]:
""" Randomly generate one sequence batch on the device.
Returns:
A sequence batch, List[Tensor]. Each element of the returned list is a
2D tensor with a shape of [sequence_length, input_dim].
"""
random.seed(1234)
torch.manual_seed(1234)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(1234)
seq_len = [random.randint(min_len, max_len) for _ in range(batch_size)]
batch = torch.randn(sum(seq_len), input_size, device=device)
offset = 0
batch_list = []
for i in range(batch_size):
a_seq = torch.as_strided(batch,
size=(seq_len[i], input_size),
stride=(input_size, 1),
storage_offset=offset)
offset += seq_len[i] * input_size
batch_list.append(a_seq)
return batch_list, seq_len
def test_mismatching_stride_no_check(self, device):
a = torch.rand((2, 2), device=device)
b = torch.as_strided(a.clone().t().contiguous(), a.shape,
a.stride()[::-1])
for fn in self.assert_fns():
fn(a, b, check_stride=False)
def test_mismatching_stride_no_check(self, device):
actual = torch.rand((2, 2), device=device)
expected = torch.as_strided(actual.clone().t().contiguous(),
actual.shape,
actual.stride()[::-1])
for fn in self.assert_fns_with_inputs(actual, expected):
fn(check_stride=False)
def test_mismatching_stride(self):
actual = torch.empty((2, 2))
expected = torch.as_strided(actual.clone().t().contiguous(), actual.shape, actual.stride()[::-1])
for fn in assert_fns_with_inputs(actual, expected):
with self.assertRaisesRegex(AssertionError, "stride"):
fn()
def __call__(self, x: torch.Tensor) -> torch.Tensor:
"""Returns the :py:class:`torch.Tensor` after adding context frames.
Args:
x: :py:class:`torch.Tensor` with size ``(1, features, seq_len)``.
Returns:
A :py:class:`torch.Tensor` with size ``(2*n_context + 1, features,
seq_len)``.
"""
# Pad to ensure first and last n_context frames in original sequence
# have at least n_context frames to their left and right respectively.
assert x.size(0) == 1
x = x.squeeze().T
steps, features = x.shape
padding = torch.zeros((self.n_context, features), dtype=x.dtype)
x = torch.cat((padding, x, padding))
window_size = self.n_context + 1 + self.n_context
strides = x.stride()
strided_x = torch.as_strided(
x,
# Shape of the new array.
(steps, window_size, features),
# Strides of the new array (bytes to step in each dim).
(strides[0], strides[0], strides[1]),
)
return strided_x.clone().detach().permute(1, 2, 0)
def _pool1d(tensor, kernel_size: int = 2, stride: int = 2, mode="max"):
output_shape = (
(tensor.shape[0] - kernel_size) // stride + 1,
)
kernel_size = (kernel_size,)
# kernel_size = (kernel_size, kernel_size)
b = torch.ones(tensor.shape) # when torch.Tensor.stride() is supported: replace with A.stride()
a_strides = b.stride()
a_w = torch.as_strided(
tensor,
size=output_shape + kernel_size,
stride=(stride * a_strides[0],) + a_strides,
)
a_w = a_w.reshape(-1, *kernel_size)
result = []
if mode == "max":
for channel in range(a_w.shape[0]):
result.append(a_w[channel].max())
elif mode == "mean":
for channel in range(a_w.shape[0]):
result.append(torch.mean(a_w[channel]))
else:
raise ValueError("unknown pooling mode")
result = torch.stack(result).reshape(output_shape)
return result
def test_mismatching_stride(self, device):
a = torch.empty((2, 2), device=device)
b = torch.as_strided(a.clone().t().contiguous(), a.shape,
a.stride()[::-1])
for fn in self.assert_fns():
with self.assertRaisesRegex(AssertionError, "stride"):
fn(a, b)
def max_pool2d_layer_numpy(self, x):
b, c, h, w = x.shape
b_strided, c_strided, h_strided, w_strided = x.stride()
x_strided = as_strided(x, (b, c, h // 2, w // 2, 2, 2),
(b_strided, c_strided, h_strided * 2,
w_strided * 2, h_strided, w_strided))
result = torch.amax(x_strided, dim=(-1, -2))
return result
def __call__(self, feat_data, item):
sliced = super(AsFramedSlice, self).__call__(feat_data, item)
if self.as_strided:
if isinstance(sliced, np.ndarray):
as_strided = lambda tensor: torch.as_strided(
torch.from_numpy(tensor),
size=(self.length - self.frame_size + 1, self.frame_size),
stride=(1, 1))
else:
as_strided = lambda tensor: torch.as_strided(
tensor,
size=(self.length - self.frame_size + 1, self.frame_size),
stride=(1, 1))
with torch.no_grad():
return as_strided(sliced)
else:
return sliced.reshape(-1, self.frame_size)
def signal_framing(signal, frame_length, frame_step):
shape = list(signal.size())
shape = shape[:-1] + [
(shape[-1] - frame_length + frame_step) // frame_step, frame_length
]
strides = list(signal.stride())
strides.insert(-1, frame_step * strides[-1])
signal = torch.as_strided(signal, size=shape, stride=strides)
return signal
def unfold(tensor, size, step, dilation=1):
assert tensor.dim() == 1
o_stride = tensor.stride(0)
numel = tensor.numel()
new_stride = (step * o_stride, dilation * o_stride)
new_size = ((numel - (dilation * (size - 1) + 1)) // step + 1, size)
if new_size[0] < 1:
new_size = (0, size)
return torch.as_strided(tensor, new_size, new_stride)
def forward(self, y, y_hat):
"""Calculate time domain loss
Args:
y (Tensor): real waveform
y_hat (Tensor): fake waveform
Return:
total_loss (Tensor): total loss of time domain
"""
# Energy loss & Time loss & Phase loss
loss_e = torch.zeros(self.len).to(y)
loss_t = torch.zeros(self.len).to(y)
loss_p = torch.zeros(self.len).to(y)
for i in range(self.len):
y_tmp = torch.as_strided(y, self.shapes[i], self.strides[i])
y_hat_tmp = torch.as_strided(y_hat, self.shapes[i],
self.strides[i])
loss_e[i] = F.l1_loss(torch.mean(y_tmp**2, dim=-1),
torch.mean(y_hat_tmp**2, dim=-1))
loss_t[i] = F.l1_loss(torch.mean(y_tmp, dim=-1),
torch.mean(y_hat_tmp, dim=-1))
if i == 0:
y_phase = F.pad(y_tmp.transpose(1, 2),
(1, 0), "constant", 0) - F.pad(
y_tmp.transpose(1, 2),
(0, 1), "constant", 0)
y_hat_phase = F.pad(y_hat_tmp.transpose(1, 2),
(1, 0), "constant", 0) - F.pad(
y_hat_tmp.transpose(1, 2),
(0, 1), "constant", 0)
else:
y_phase = F.pad(y_tmp, (1, 0), "constant", 0) - F.pad(
y_tmp, (0, 1), "constant", 0)
y_hat_phase = F.pad(y_hat_tmp, (1, 0), "constant", 0) - F.pad(
y_hat_tmp, (0, 1), "constant", 0)
loss_p[i] = F.l1_loss(y_phase, y_hat_phase)
total_loss = torch.sum(loss_e) + torch.sum(loss_t) + torch.sum(loss_p)
return total_loss
def _create_views(self, tensor):
views = []
for stride in self.strides:
outdim = tensor.size(0) - stride + 1
view = torch.as_strided(tensor, (outdim, stride, self.dim),
(self.dim, self.dim, 1))
views.append(view)
return views
def unfold(tensor, size, step, dilation=1):
assert tensor.dim() == 1
o_stride = tensor.stride(0)
numel = tensor.numel()
new_stride = (step * o_stride, dilation * o_stride)
new_size = ((numel - (dilation * (size - 1) + 1)) // step + 1, size)
if new_size[0] < 1:
#new_size = (0, size)
#not to exclude video shorther than size
return tensor.unsqueeze(0)
return torch.as_strided(tensor, new_size, new_stride)
def window_view(x, dim, size, stride, at=None):
r"""Returns a sliding window view into the tensor.
Similar to `torch.unfold()`, but the window dimensions of size `size`
is placed right after `dim` (by default), and not appended.
Arguments
---------
x : torch.tensor
Time series of measurement values
dim : int
Axis along which the periodogram is computed, i.e. ``dim=-1``.
size : int
The size of the sliding windows.
stride : int
The step between two sliding windows.
at : int, optional
The dimension at which to put the slice of each window.
Returns
-------
x_view : torch.tensor
The view into a sliding window. The returned tensor is not, in
general, contiguous in memory.
"""
if size <= 0:
raise ValueError(f"""`size` must be a positive integer.""")
if stride < 0:
raise ValueError(f"""`stride` must be a nonnegative integer.""")
dim = fix_dim(dim, x.dim())
if x.shape[dim] < size:
raise ValueError(f"""`x` at dim {dim} is too short ({x.shape[dim]}) """
f"""for this window size ({size}).""")
if at is None:
at = dim + 1
at = fix_dim(at, x.dim() + 1)
# compute new shape and strides
shape, strides = list(x.size()), list(x.stride())
strided_size = ((shape[dim] - size + 1) + stride - 1) // stride
# new shape and stride structure
shape_view = shape[:dim] + [strided_size] + shape[dim+1:]
shape_view.insert(at, size)
strides_view = strides[:dim] + [strides[dim] * stride] + strides[dim+1:]
strides_view.insert(at, strides[dim])
# differentiable strided view
return torch.as_strided(x, shape_view, strides_view)
def vjp(outputs, inputs, **kwargs):
if torch.is_tensor(inputs):
inputs = [inputs]
_dummy_inputs = [torch.as_strided(i, (), ())
for i in inputs] # Workaround for PyTorch bug #39784.
if torch.is_tensor(outputs):
outputs = [outputs]
outputs = make_seq_requires_grad(outputs)
_vjp = torch.autograd.grad(outputs, inputs, **kwargs)
return convert_none_to_zeros(_vjp, inputs)
def forward(self, y, y_hat):
"""Calculate time domain loss
Args:
y (Tensor): real waveform
y_hat (Tensor): fake waveform
Return:
total_loss (Tensor): total loss of time domain
"""
# Energy loss
loss_e = torch.zeros(self.len).to(y)
for i in range(self.len):
y_energy = torch.as_strided(y**2, self.shapes[i], self.strides[i])
y_hat_energy = torch.as_strided(y_hat**2, self.shapes[i],
self.strides[i])
loss_e[i] = F.l1_loss(torch.mean(y_energy, dim=-1),
torch.mean(y_hat_energy, dim=-1))
# Time loss
loss_t = torch.zeros(self.len).to(y)
for i in range(self.len):
y_time = torch.as_strided(y, self.shapes[i], self.strides[i])
y_hat_time = torch.as_strided(y_hat, self.shapes[i],
self.strides[i])
loss_t[i] = F.l1_loss(torch.mean(y_time, dim=-1),
torch.mean(y_hat_time, dim=-1))
# Phase loss
y_phase = F.pad(y,
(1, 0), "constant", 0) - F.pad(y,
(0, 1), "constant", 0)
y_hat_phase = F.pad(y_hat, (1, 0), "constant", 0) - F.pad(
y_hat, (0, 1), "constant", 0)
loss_p = F.l1_loss(y_phase, y_hat_phase)
total_loss = torch.sum(loss_e) + torch.sum(loss_t) + loss_p
return total_loss
def tensor_creation_ops(self):
i = torch.tensor([[0, 1, 1], [2, 0, 2]])
v = torch.tensor([3, 4, 5], dtype=torch.float32)
real = torch.tensor([1, 2], dtype=torch.float32)
imag = torch.tensor([3, 4], dtype=torch.float32)
inp = torch.tensor([-1.5, 0.0, 2.0])
values = torch.tensor([0.5])
quantized = torch.quantize_per_channel(
torch.tensor([[-1.0, 0.0], [1.0, 2.0]]),
torch.tensor([0.1, 0.01]),
torch.tensor([10, 0]),
0,
torch.quint8,
)
return (
torch.tensor([[0.1, 1.2], [2.2, 3.1], [4.9, 5.2]]),
# torch.sparse_coo_tensor(i, v, [2, 3]), # not work for iOS
torch.as_tensor([1, 2, 3]),
torch.as_strided(torch.randn(3, 3), (2, 2), (1, 2)),
torch.zeros(2, 3),
torch.zeros((2, 3)),
torch.zeros([2, 3], out=i),
torch.zeros(5),
torch.zeros_like(torch.empty(2, 3)),
torch.ones(2, 3),
torch.ones((2, 3)),
torch.ones([2, 3]),
torch.ones(5),
torch.ones_like(torch.empty(2, 3)),
torch.arange(5),
torch.arange(1, 4),
torch.arange(1, 2.5, 0.5),
torch.range(1, 4),
torch.range(1, 4, 0.5),
torch.linspace(3.0, 3.0, steps=1),
torch.logspace(start=2, end=2, steps=1, base=2.0),
torch.eye(3),
torch.empty(2, 3),
torch.empty_like(torch.empty(2, 3), dtype=torch.int64),
torch.empty_strided((2, 3), (1, 2)),
torch.full((2, 3), 3.141592),
torch.full_like(torch.full((2, 3), 3.141592), 2.71828),
torch.quantize_per_tensor(
torch.tensor([-1.0, 0.0, 1.0, 2.0]), 0.1, 10, torch.quint8
),
torch.dequantize(quantized),
torch.complex(real, imag),
torch.polar(real, imag),
torch.heaviside(inp, values),
)
def get_frame(audio, step_size, center):
if center:
audio = nn.functional.pad(audio, pad=(512, 512))
# make 1024-sample frames of the audio with hop length of 10 milliseconds
hop_length = int(16000 * step_size / 1000)
n_frames = 1 + int((len(audio) - 1024) / hop_length)
assert audio.dtype == torch.float32
itemsize = 1 # float32 byte size
frames = torch.as_strided(audio, size=(1024, n_frames), stride=(itemsize, hop_length * itemsize))
frames = frames.transpose(0, 1).clone()
frames -= (torch.mean(frames, axis=1).unsqueeze(-1))
frames /= (torch.std(frames, axis=1).unsqueeze(-1))
return frames
def convolution(self, m):
f = self.weight
m_h, m_w = m.shape[-2:]
f_h, f_w = f.shape[-2:]
batch_size = m.shape[0]
m_c = m.shape[1]
Hout = m_h - f_h + 1
Wout = m_w - f_w + 1
stride_batch_size, stride_c, stride_h, stride_w = m.stride()
m_strided = as_strided(m, (batch_size, Hout, Wout, m_c, f_h, f_w),
(stride_batch_size, stride_h, stride_w,
stride_c, stride_h, stride_w))
result = einsum('bmncuv,kcuv->bkmn', m_strided, f)
return result
def test_stride(self):
cpu_ones = torch.ones(3, 3)
ort_ones = cpu_ones.to("ort")
y = torch.as_strided(ort_ones, (2, 2), (1, 2))
assert y.size() == (2, 2)
assert y.is_contiguous() == False
contiguous_y = y.contiguous()
w = torch.ones((2, 3))
ort_w = w.to("ort")
z = torch.zeros((2, 3))
ort_z = z.to("ort")
ort_z = torch.addmm(ort_z, contiguous_y, ort_w)
cpu_z = torch.addmm(z, torch.ones(2, 2), w)
assert torch.allclose(ort_z.cpu(), cpu_z)
def preprocess_audio_batch(audio, sr, center=True, hop_size=0.1, sampler="julian"):
if audio.ndim == 3:
audio = torch.mean(audio, axis=2)
if sr != TARGET_SR:
if sampler == "julian":
audio = julius.resample_frac(audio, sr, TARGET_SR)
elif sampler == "resampy":
audio = torch.tensor(
resampy.resample(
audio.detach().cpu().numpy(),
sr_orig=sr,
sr_new=TARGET_SR,
filter="kaiser_best",
),
dtype=audio.dtype,
device=audio.device,
)
else:
raise ValueError("Only julian and resampy works!")
frame_len = TARGET_SR
hop_len = int(hop_size * TARGET_SR)
if center:
audio = center_audio(audio, frame_len)
audio = pad_audio(audio, frame_len, hop_len)
n_frames = 1 + int((audio.size()[1] - frame_len) / float(hop_len))
x = []
xframes_shape = None
for i in range(audio.shape[0]):
xframes = (
torch.as_strided(
audio[i],
size=(frame_len, n_frames),
stride=(1, hop_len),
)
.transpose(0, 1)
.unsqueeze(1)
)
if xframes_shape is None:
xframes_shape = xframes.shape
assert xframes.shape == xframes_shape
x.append(xframes)
x = torch.vstack(x)
return x
def unfold(tensor, frames_per_clip, frames_between_clips, dilation=1):
"""
similar to tensor.unfold, but with the dilation
and specialized for 1d tensors
Returns all consecutive windows of `size` elements, with
`step` between windows. The distance between each element
in a window is given by `dilation`.
"""
assert tensor.dim() == 1
o_stride = tensor.stride(0)
numel = tensor.numel()
new_stride = (frames_between_clips * o_stride, dilation * o_stride)
new_size = ((numel - (dilation * (frames_per_clip - 1) + 1)) // frames_between_clips + 1, frames_per_clip)
if new_size[0] < 1:
new_size = (0, frames_per_clip)
return torch.as_strided(tensor, new_size, new_stride)
def t_window(x, size, shift=None, stride=1):
try:
nd = len(size)
except TypeError:
size = tuple(size for i in x.shape)
nd = len(size)
if nd != x.ndimension():
raise ValueError("size has length {0} instead of "
"x.ndim which is {1}".format(len(size),
x.ndimension()))
out_shape = tuple(xi - wi + 1 for xi, wi in zip(x.shape, size)) + size
if not all(i > 0 for i in out_shape):
raise ValueError("size is bigger than input array along at "
"least one dimension")
out_strides = x.stride() * 2
return t.as_strided(x, out_shape, out_strides)
def batch_stripe(a):
"""
Get a diagonal stripe of a matrix m x n, where n > m
this implementation also takes into account batched matrices,
so the stripe is calculated over a batch x for a matrix of size[x, m, n]
"""
# another solution
# a = a[::-1] # ValueError: negative step not yet supported
# do the usual left top to right bottom
# return a[::-1]
b, i, j = a.size()
assert i > j
b_s, k, l = a.stride()
# left top to right bottom
return torch.as_strided(a, (b, i - j + 1, j), (b_s, k, k + l))
def unfold(tensor: torch.Tensor,
size: int,
step: int,
dilation: int = 1) -> torch.Tensor:
"""
similar to tensor.unfold, but with the dilation
and specialized for 1d tensors
Returns all consecutive windows of `size` elements, with
`step` between windows. The distance between each element
in a window is given by `dilation`.
"""
assert tensor.dim() == 1
o_stride = tensor.stride(0)
numel = tensor.numel()
new_stride = (step * o_stride, dilation * o_stride)
new_size = ((numel - (dilation * (size - 1) + 1)) // step + 1, size)
if new_size[0] < 1:
new_size = (0, size)
return torch.as_strided(tensor, new_size, new_stride)
def jvp(outputs, inputs, grad_inputs=None, **kwargs):
# Unlike `torch.autograd.functional.jvp`, this function avoids repeating forward computation.
if torch.is_tensor(inputs):
inputs = [inputs]
_dummy_inputs = [torch.as_strided(i, (), ())
for i in inputs] # Workaround for PyTorch bug #39784.
if torch.is_tensor(outputs):
outputs = [outputs]
outputs = make_seq_requires_grad(outputs)
dummy_outputs = [torch.zeros_like(o, requires_grad=True) for o in outputs]
_vjp = torch.autograd.grad(outputs,
inputs,
grad_outputs=dummy_outputs,
create_graph=True,
allow_unused=True)
_vjp = make_seq_requires_grad(convert_none_to_zeros(_vjp, inputs))
_jvp = torch.autograd.grad(_vjp,
dummy_outputs,
grad_outputs=grad_inputs,
**kwargs)
return convert_none_to_zeros(_jvp, dummy_outputs)
