def nllloss_1d(self, input, target):
n = input.shape[0]
idx = flow.unsqueeze(flow.arange(0, n, 1), dim=1)
target = flow.unsqueeze(target, dim=1)
t = flow.cat([idx, target], dim=1)
res = self._gather_nd_op(input, t)[0]
return res
def _transform_input(self, x: Tensor) -> Tensor:
if self.transform_input:
x_ch0 = flow.unsqueeze(x[:, 0],
1) * (0.229 / 0.5) + (0.485 - 0.5) / 0.5
x_ch1 = flow.unsqueeze(x[:, 1],
1) * (0.224 / 0.5) + (0.456 - 0.5) / 0.5
x_ch2 = flow.unsqueeze(x[:, 2],
1) * (0.225 / 0.5) + (0.406 - 0.5) / 0.5
x = flow.cat((x_ch0, x_ch1, x_ch2), 1)
return x
def _test_unsqueeze_different_dim(test_case, device):
np_arr = np.random.rand(4, 5, 6, 7)
x = flow.tensor(np_arr, dtype=flow.float32, device=flow.device(device))
for axis in range(-5, 5):
y = flow.unsqueeze(x, dim=axis)
output = np.expand_dims(np_arr, axis=axis)
test_case.assertTrue(np.allclose(output, y.numpy(), 1e-05, 1e-05))
def forward(self, x):
if x.dim() >= 3:
raise RuntimeError(
"{} accept 1/2D tensor as input, but got {:d}".format(
self.__name__, x.dim()
)
)
# when inference, only one utt
if x.dim() == 1:
x = flow.unsqueeze(x, 0)
# n x 1 x S => n x N x T
w = F.relu(self.encoder_1d(x))
# n x B x T
y = self.proj(self.ln(w))
# n x B x T
y = self.repeats(y)
# n x 2N x T
e = flow.chunk(self.mask(y), self.num_spks, 1)
# n x N x T
if self.non_linear_type == "softmax":
m = self.non_linear(flow.stack(e, dim=0), dim=0)
else:
m = self.non_linear(flow.stack(e, dim=0))
# spks x [n x N x T]
s = [w * m[n] for n in range(self.num_spks)]
# spks x n x S
return [self.decoder_1d(x, squeeze=True) for x in s]
def test_unsqueeze_different_dim(test_case):
np_arr = np.random.rand(4, 5, 6, 7)
x = flow.Tensor(np_arr)
for axis in range(-5, 5):
y = flow.unsqueeze(x, dim=axis)
output = np.expand_dims(np_arr, axis=axis)
test_case.assertTrue(np.allclose(output, y.numpy(), rtol=1e-05))
def _prob_in_top_k(
self, clean_values, noisy_values, noise_stddev, noisy_top_values
):
"""Helper function to NoisyTopKGating.
Computes the probability that value is in top k, given different random noise.
This gives us a way of backpropagating from a loss that balances the number
of times each expert is in the top k experts per example.
In the case of no noise, pass in None for noise_stddev, and the result will
not be differentiable.
Args:
clean_values: a `Tensor` of shape [batch, n].
noisy_values: a `Tensor` of shape [batch, n]. Equal to clean values plus
normally distributed noise with standard deviation noise_stddev.
noise_stddev: a `Tensor` of shape [batch, n], or None
noisy_top_values: a `Tensor` of shape [batch, m].
"values" Output of tf.top_k(noisy_top_values, m). m >= k+1
Returns:
a `Tensor` of shape [batch, n].
"""
batch = clean_values.size(0)
m = noisy_top_values.size(1)
top_values_flat = noisy_top_values.flatten()
threshold_positions_if_in = (
flow.arange(batch, device=noisy_values.device) * m + self.k
)
threshold_if_in = flow.unsqueeze(
flow.gather(top_values_flat, 0, threshold_positions_if_in), 1
)
is_in = flow.gt(noisy_values, threshold_if_in)
threshold_positions_if_out = threshold_positions_if_in - 1
threshold_if_out = flow.unsqueeze(
flow.gather(top_values_flat, 0, threshold_positions_if_out), 1
)
# is each value currently in the top k.
prob_if_in = cdf((clean_values - threshold_if_in) / noise_stddev)
prob_if_out = cdf((clean_values - threshold_if_out) / noise_stddev)
prob = flow.where(is_in, prob_if_in, prob_if_out)
return prob
def _test_unsqueeze_backward(test_case, device):
np_arr = np.random.rand(2, 3, 4, 5)
x = flow.tensor(np_arr,
dtype=flow.float32,
device=flow.device(device),
requires_grad=True)
y = flow.unsqueeze(x, dim=1).sum()
y.backward()
test_case.assertTrue(
np.allclose(x.grad.numpy(), np.ones((2, 3, 4, 5)), 1e-05, 1e-05))
def forward(self, x, squeeze=False):
"""
x: N x L or N x C x L
"""
if x.dim() not in [2, 3]:
raise RuntimeError("{} accept 2/3D tensor as input".format(self.__name__))
x = super().forward(x if x.dim() == 3 else flow.unsqueeze(x, 1))
if squeeze:
x = flow.squeeze(x)
return x
def __getitem__(self, idx):
p = self.files[idx]
filename = os.path.basename(p)
speaker = filename.split(sep="_", maxsplit=1)[0]
label = self.encoder.transform([speaker])[0]
mcep = np.load(p)
mcep = flow.Tensor(mcep)
mcep = flow.unsqueeze(mcep, 0)
return (
mcep,
flow.tensor(speakers.index(speaker), dtype=flow.long),
flow.tensor(label, dtype=flow.float),
)
def _test_unsqueeze(test_case, device):
np_arr = np.random.rand(2, 6, 9, 3)
x = flow.tensor(np_arr, dtype=flow.float32, device=flow.device(device))
y = flow.unsqueeze(x, dim=1)
output = np.expand_dims(np_arr, axis=1)
test_case.assertTrue(np.allclose(output, y.numpy(), 1e-05, 1e-05))
def test_dynamic_attrs(test_case):
x = flow.full((2, 3), 3.0)
y = flow.unsqueeze(x, dim=1)
test_case.assertEqual(y.shape, flow.Size((2, 1, 3)))
y = flow.unsqueeze(x, dim=2)
test_case.assertEqual(y.shape, flow.Size((2, 3, 1)))
def test_unsqueeze(test_case):
np_arr = np.random.rand(2, 6, 9, 3)
x = flow.Tensor(np_arr)
y = flow.unsqueeze(x, dim=1)
output = np.expand_dims(np_arr, axis=1)
test_case.assertTrue(np.allclose(output, y.numpy(), rtol=1e-05))