def _TestStreamStepHelper(self, **kwargs):
"""Main helper method."""
batch_size, max_seqlen, input_dim = 2, 32, kwargs['input_dim']
stride = kwargs.get('stride', 1)
# max_seqlen is divisible by stride.
assert max_seqlen % stride == 0
right_context = kwargs.get('right_context', 0)
# Prepares inputs.
inputs, paddings = self._GetInputs(batch_size, max_seqlen, input_dim)
# Gets params
p = self._GetParams(**kwargs)
# Builds graph.
with self.session(use_gpu=False) as sess:
l = p.Instantiate()
init_op = tf.global_variables_initializer()
fprop_out = self._FProp(l, inputs, paddings)
base_outputs = self._GetFPropOutput(fprop_out)
out_rank = py_utils.GetRank(base_outputs)
base_outputs *= py_utils.AppendDims(1. - paddings, out_rank - 2)
try:
state = l.zero_state(batch_size)
except TypeError:
state = l.zero_state(l.theta, batch_size)
outputs = []
for i in range(max_seqlen // stride +
int(math.ceil(right_context / stride))):
if i < max_seqlen // stride:
step_inputs = inputs[:, stride * i:stride * (i + 1)]
step_paddings = paddings[:, stride * i:stride * (i + 1)]
else:
step_inputs = tf.zeros_like(inputs[:, 0:stride])
step_paddings = tf.ones_like(paddings[:, 0:stride])
output, _, state = l.StreamStep(l.theta, step_inputs, step_paddings,
state)
outputs.append(output)
outputs = tf.concat(outputs, axis=1)
outputs = self._NormalizeStreamStepOutput(outputs, paddings,
right_context, max_seqlen)
sess.run(init_op)
expected, actual = sess.run([base_outputs, outputs])
print(f'expected: {repr(expected)}, {expected.shape}')
print(f'actual: {repr(actual)}, {actual.shape}')
print(f'np.sum(np.abs(expected)): {np.sum(np.abs(expected))}')
print(f'np.sum(np.abs(actual)): {np.sum(np.abs(actual))}')
tol = kwargs.get('tol', 1e-6)
self.assertAllClose(expected, actual, atol=tol, rtol=tol)
def _NormalizeStreamStepOutput(self,
outputs,
paddings,
right_context,
max_seqlen,
num_layers=1):
# outputs has right_context * num_layers-frames delay from inputs.
outputs = outputs[:, right_context * num_layers:]
# later outputs corresponds to padded inputs to complete the last frame's
# right context.
outputs = outputs[:, :max_seqlen]
out_rank = py_utils.GetRank(outputs)
paddings = paddings[:, :max_seqlen]
return outputs * py_utils.AppendDims(1. - paddings, out_rank - 2)
def StreamStep(self, theta, inputs, paddings, state0):
"""Apply a single step of convolution to input_tensor.
Only supports 1d causal convolution. Doesn't support dilation.
Args:
theta: A NestedMap of layer params.
inputs: A Tensor of shape [b, t, 1, c]
paddings: A 0/1 valued tensor of shape [b, t].
state0: A NestedMap of tensors of the same struct as returned by
zero_state().
Returns:
outputs: A Tensor of shape [b, t, 1, c * channel_multiplier]
padding: the same as input paddings.
state1: A NestedMap of the same struct as input state
"""
p = self.params
assert p.filter_shape[1] == 1, (
'StreamStep only supports 1d causal convolution.')
assert p.filter_stride[0] == 1, (
'StreamStep doesn\'t support striding')
assert p.dilation_rate == (1,
1), ('StreamStep doesn\'t support dilation')
with tf.name_scope(p.name):
inputs = py_utils.HasShape(inputs, [-1, -1, 1, p.filter_shape[2]])
paddings = py_utils.HasShape(paddings,
py_utils.GetShape(inputs)[:2])
q = py_utils.GetShape(paddings)[1]
padded_inputs = py_utils.ApplyPadding(
py_utils.AppendDims(paddings, 2), inputs)
concat_inputs = tf.concat([state0.context, padded_inputs], axis=1)
outputs = tf.nn.depthwise_conv2d(concat_inputs,
self._GetWeight(theta),
strides=(1, 1, 1, 1),
dilations=(1, 1),
data_format='NHWC',
padding='VALID')
if p.bias:
outputs = tf.nn.bias_add(outputs, theta.b)
new_context = concat_inputs[:, q:]
return outputs, paddings, py_utils.NestedMap(context=new_context)
def StreamStep(self, theta, inputs, paddings, state0):
"""Apply a singele step of convolution to input_tensor.
Only supports 1d causal convolution. Doesn't support dilation.
Args:
theta: A NestedMap of layer params.
inputs: A Tensor of shape [b, t, 1, c]
paddings: A 0/1 valued tensor of shape [b, t].
state0: A NestedMap of tensors of the same struct as returned by
zero_state().
Returns:
outputs: A Tensor of shape [b, t, 1, c]
padding: the same as input paddings.
state1: A NestedMap of the same struct as input state
"""
p = self.params
assert p.filter_shape[1] == 1, (
'StreamStep only supports 1d causal convolution.')
assert all(stride == 1 for stride in p.filter_stride), (
f'StreamStep doesn\'t support striding: {p.filter_stride}')
assert p.dilation_rate == (1,
1), ('StreamStep doesn\'t support dilation')
with tf.name_scope(p.name):
inputs = py_utils.HasShape(inputs, [-1, -1, 1, p.filter_shape[2]])
paddings = py_utils.HasShape(paddings,
py_utils.GetShape(inputs)[:2])
concat_inputs = tf.concat([
state0.context, inputs * (1 - py_utils.AppendDims(paddings, 2))
],
axis=1)
outputs = tf.nn.conv2d(concat_inputs,
self._GetWeight(theta),
strides=p.filter_stride,
dilations=p.dilation_rate,
data_format='NHWC',
padding='VALID')
new_context = concat_inputs[:, -(p.filter_shape[0] - 1):]
return outputs, paddings, py_utils.NestedMap(context=new_context)
def expand_pad(pad): # pylint:disable=invalid-name return py_utils.AppendDims(pad, 2)