def __call__(self, in_obj):
in_axes = in_obj.axes
if in_axes.channel_axis() is None:
red_axes = ng.make_axes(in_axes.recurrent_axis()) + in_axes.batch_axes()
else:
red_axes = in_axes - in_axes.channel_axis()
out_axes = in_axes - red_axes
in_obj = ng.flatten(in_obj, out_axes | red_axes.flatten(force=True))
if self.gamma is None:
self.gvar = self.gvar or ng.persistent_tensor(axes=out_axes, initial_value=1.0)
self.gmean = self.gmean or ng.persistent_tensor(axes=out_axes, initial_value=0.0)
self.gamma = ng.variable(axes=out_axes,
initial_value=self.init_gamma,
scope=self.scope).named('gamma')
self.beta = ng.variable(axes=out_axes,
initial_value=self.init_beta,
scope=self.scope).named('beta')
xmean = ng.mean(in_obj, out_axes=out_axes)
xvar = ng.variance(in_obj, out_axes=out_axes)
if Layer.inference_mode:
return ng.unflatten(self.gamma * ((in_obj - self.gmean) *
ng.reciprocal(ng.sqrt(self.gvar + self.eps))) + self.beta)
else:
return ng.sequential([
ng.assign(self.gmean, self.gmean * self.rho + xmean * (1.0 - self.rho)),
ng.assign(self.gvar, self.gvar * self.rho + xvar * (1.0 - self.rho)),
ng.unflatten(self.gamma * ((in_obj - xmean) *
ng.reciprocal(ng.sqrt(xvar + self.eps))) + self.beta)
])
def train_outputs(self, in_obj):
"""
Arguments:
in_obj (Tensor): object that provides the lookup indices
"""
in_obj.axes.find_by_short_name('time')[0].add_role(ar.time)
in_obj.axes.find_by_short_name('time')[0].is_recurrent = True
in_obj = ng.axes_with_role_order(in_obj, self.role_order)
in_obj = ng.flatten(in_obj)
in_axes = in_obj.axes
self.lut_v_axis = ng.make_axis(self.vocab_size).named('V')
self.lut_f_axis = ng.make_axis(self.embed_dim).named('F')
self.w_axes = ng.make_axes([self.lut_v_axis, self.lut_f_axis])
self.lut_o_axes = in_axes + ng.make_axes([self.lut_f_axis])
self.o_axes = ng.make_axes([self.lut_f_axis]) + in_axes[0].axes
self.W = ng.variable(axes=self.w_axes,
initial_value=self.lut_init(
self.w_axes, self.lut_v_axis,
self.pad_idx)).named('W')
lut_result = ng.lookuptable(self.W,
in_obj,
self.lut_o_axes,
update=self.update,
pad_idx=self.pad_idx)
return ng.axes_with_order(ng.unflatten(lut_result), self.o_axes)
def test_lut(lut_args):
"""
test lut fprop and bprop
"""
pad_idx = 0
with ExecutorFactory() as ex:
vocab_size, embed_dim, bsz, seq_len, mem_size = lut_args
V = ng.make_axis(vocab_size)
F = ng.make_axis(embed_dim)
M = ng.make_axis(mem_size)
ax.N.length = bsz
ax.REC.length = seq_len
# Multi-axis input to LUT
ax_idx = ng.make_axes([M, ax.REC, ax.N])
ax_lut = ng.make_axes([V, F])
lut = ng.placeholder(ax_lut)
idx = ng.placeholder(ax_idx)
idx_flat = ng.flatten(idx)
ax_out = idx_flat.axes | ng.make_axes([F])
# fprop
lut_out_ng = ng.lookuptable(lut, idx_flat, ax_out, pad_idx=pad_idx)
fprop_fun = ex.executor(lut_out_ng, lut, idx)
# bprop
update_error = ng.placeholder(ax_out)
update_out_ng = lookuptable_update(update_error, lut, idx, lut_out_ng)
update_fun = ex.executor(update_out_ng, update_error, lut, idx)
# provide actual inputs and execute the graph
lut_value = rng.uniform(-1, 1, lut.axes)
idx_value = rng.random_integers(0, vocab_size - 1, idx.axes)
fprop_lut = fprop_fun(lut_value, idx_value).copy()
# compare fprop
fprop_ref = lut_fprop_ref(lut_value, idx_value)
ng.testing.assert_allclose(fprop_lut, fprop_ref, rtol=0.0, atol=1.0e-5)
# provide actual delta and execute the update op
update_value = rng.uniform(-1, 1, update_error.axes)
update_lut = update_fun(update_value, lut_value, idx_value).copy()
# compare bprop (udpate)
update_ref = lut_update_ref(update_value,
lut_value,
idx_value,
pad_idx=pad_idx)
ng.testing.assert_allclose(update_lut,
update_ref,
rtol=0.0,
atol=1.0e-5)
def Reshape(self, tf_node, inputs):
"""
Reshapes a tensor.
Arguments:
tf_node: NodeDef object, the tensorflow node to convert.
inputs: List of ngraph Ops as inputs to this node.
Returns:
A ngraph Op corresponding to the tensorflow node.
Inputs to tf_node:
tensor, shape, name
"""
# TODO: currently only support constants and flatten to 1d and 2d
# get inputs
tensor, shape = inputs
def get_flatten_idx(shape_i, shape_o):
"""
check if flattening shape is valid
Args:
shape_i: input tensor shape
shape_o: output flattend tensor shape
Returns:
None if flatten not valid, otherwise the flatten_at index
"""
return None
# get input and output shape
shape_i = tensor.shape.lengths
shape_o = tuple(shape.const.astype(int))
if np.prod(shape_i) != np.prod(shape_o):
raise ValueError("Total size of input and output dimension "
"mismatch.")
if tensor.const is not None:
# reshape const
np_val = np.reshape(tensor.const, shape_o)
return ng.constant(np_val,
make_pos_axes(np_val.shape)).named(tf_node.name)
else:
ndims_o = len(shape_o)
if ndims_o != 1 and ndims_o != 2:
raise NotImplementedError("Reshape can only support flatten"
"to 1d or 2d.")
if ndims_o == 1:
tensor = ng.flatten(tensor)
else:
cumprods = list(np.cumprod(shape_i))
flatten_at_idx = cumprods.index(shape_o[0]) + 1
tensor = ng.flatten_at(tensor, flatten_at_idx)
res = ng.cast_axes(tensor, make_pos_axes(shape_o))
return res.named(tf_node.name)
def test_lut(lut_args):
"""
test lut fprop and bprop
"""
pad_idx = 0
with ExecutorFactory() as ex:
vocab_size, embed_dim, bsz, seq_len, mem_size = lut_args
V = ng.make_axis(vocab_size)
F = ng.make_axis(embed_dim)
M = ng.make_axis(mem_size)
ax.N.length = bsz
ax.REC.length = seq_len
# Multi-axis input to LUT
ax_idx = ng.make_axes([M, ax.REC, ax.N])
ax_lut = ng.make_axes([V, F])
lut = ng.placeholder(ax_lut)
idx = ng.placeholder(ax_idx)
idx_flat = ng.flatten(idx)
ax_out = idx_flat.axes | ng.make_axes([F])
# fprop
lut_out_ng = ng.lookuptable(lut, idx_flat, ax_out, pad_idx=pad_idx)
fprop_fun = ex.executor(lut_out_ng, lut, idx)
# bprop
update_error = ng.placeholder(ax_out)
update_out_ng = lookuptable_update(update_error, lut, idx, lut_out_ng)
update_fun = ex.executor(update_out_ng, update_error, lut, idx)
# provide actual inputs and execute the graph
lut_value = rng.uniform(-1, 1, lut.axes)
idx_value = rng.random_integers(0, vocab_size - 1, idx.axes)
fprop_lut = fprop_fun(lut_value, idx_value).copy()
# compare fprop
fprop_ref = lut_fprop_ref(lut_value, idx_value)
ng.testing.assert_allclose(fprop_lut, fprop_ref, rtol=0.0, atol=1.0e-5)
# provide actual delta and execute the update op
update_value = rng.uniform(-1, 1, update_error.axes)
update_lut = update_fun(update_value, lut_value, idx_value).copy()
# compare bprop (udpate)
update_ref = lut_update_ref(
update_value,
lut_value,
idx_value,
pad_idx=pad_idx)
ng.testing.assert_allclose(
update_lut, update_ref, rtol=0.0, atol=1.0e-5)
def run_mini_ds2_benchmark(args, **kwargs):
device_id = kwargs.get('device_id')
inputs, train_set, eval_set = generate_ds2_data(args.max_length,
args.str_w, args.nout,
args.nbands,
args.batch_size,
args.num_iterations)
model_out = get_mini_ds2(inputs, args.nfilters, args.filter_width,
args.str_w, args.nbands, args.depth,
args.hidden_size, args.batch_norm,
args.hetr_device, device_id)
if args.bprop:
with ng.metadata(device=args.hetr_device,
device_id=device_id,
parallel=ax.N):
loss = ng.ctc(model_out, ng.flatten(inputs["char_map"]),
inputs["audio_length"], inputs["trans_length"])
optimizer = GradientDescentMomentum(learning_rate=2e-5,
momentum_coef=0.99,
gradient_clip_norm=400,
nesterov=args.nesterov)
updates = optimizer(loss)
mean_cost = ng.sequential([updates, ng.mean(loss, out_axes=())])
bprop_computation_op = ng.computation(mean_cost, "all")
benchmark = Benchmark(bprop_computation_op, train_set, inputs,
args.backend, args.hetr_device)
Benchmark.print_benchmark_results(
benchmark.time(args.num_iterations,
args.skip_iter,
'ds2_bprop',
args.visualize,
preprocess=True))
else:
fprop_computation_op = ng.computation(model_out, "all")
benchmark_fprop = Benchmark(fprop_computation_op, train_set, inputs,
args.backend, args.hetr_device)
Benchmark.print_benchmark_results(
benchmark_fprop.time(args.num_iterations,
args.skip_iter,
'ds2_fprop',
args.visualize,
preprocess=True))
def __call__(self, in_obj, **kwargs):
"""
Arguments:
in_obj (Tensor): object that provides the lookup indices
"""
LABELS = {"weight": "weight", "bias": "bias"}
in_obj = ng.axes_with_order(
in_obj,
ng.make_axes(
[in_obj.axes.recurrent_axis(),
in_obj.axes.batch_axis()]))
in_obj = ng.flatten(in_obj)
in_axes = in_obj.axes
# label lut_v_axis as shadow axis for initializers ... once #1158 is
# in, shadow axis will do more than just determine fan in/out for
# initializers.
self.lut_v_axis = ng.make_axis(self.vocab_size).named('V')
self.axes_map = shadow_axes_map([self.lut_v_axis])
self.lut_v_axis = list(self.axes_map.values())[0]
self.lut_f_axis = ng.make_axis(self.embed_dim).named('F')
self.w_axes = ng.make_axes([self.lut_v_axis, self.lut_f_axis])
self.lut_o_axes = in_axes | ng.make_axes([self.lut_f_axis])
self.o_axes = ng.make_axes([self.lut_f_axis]) | in_axes[0].axes
if not self.initialized:
self.W = ng.variable(
axes=self.w_axes,
initial_value=self.lut_init(self.w_axes, self.lut_v_axis,
self.pad_idx),
metadata={
"label": LABELS["weight"]
},
).named('LutW')
lut_result = ng.lookuptable(self.W,
in_obj,
self.lut_o_axes,
update=self.update,
pad_idx=self.pad_idx)
return ng.axes_with_order(
ng.map_roles(ng.unflatten(lut_result), self.axes_map), self.o_axes)
def __call__(self, in_obj, **kwargs):
"""
Arguments:
in_obj (Tensor): object that provides the lookup indices
"""
in_obj = ng.flatten(in_obj)
in_axes = in_obj.axes
# label lut_v_axis as shadow axis for initializers ... once #1158 is
# in, shadow axis will do more than just determine fan in/out for
# initializers.
self.lut_v_axis = ng.make_axis(self.vocab_size).named('V')
self.axes_map = shadow_axes_map([self.lut_v_axis])
self.lut_v_axis = list(self.axes_map.values())[0]
self.lut_f_axis = ng.make_axis(self.embed_dim).named('F')
self.w_axes = ng.make_axes([self.lut_v_axis, self.lut_f_axis])
self.lut_o_axes = in_axes | ng.make_axes([self.lut_f_axis])
self.o_axes = ng.make_axes([self.lut_f_axis]) | in_axes[0].axes
if not self.initialized:
self.W = ng.variable(
axes=self.w_axes,
initial_value=self.lut_init(
self.w_axes,
self.lut_v_axis,
self.pad_idx),
metadata={
"label": LABELS["weight"]},
).named('LutW')
lut_result = ng.lookuptable(
self.W,
in_obj,
self.lut_o_axes,
update=self.update,
pad_idx=self.pad_idx)
return ng.map_roles(ng.unflatten(lut_result), self.axes_map)
def reshape_workaround(
data, shape_out): # type: (TensorOp, Sequence[int]) -> TensorOp
"""Limited workaround for tensor reshape operation."""
shape_in = data.shape.lengths
if np.prod(shape_in) != np.prod(shape_out):
raise ValueError(
'Total size of input (%d) and output (%d) dimension mismatch.',
np.prod(shape_in), np.prod(shape_out))
ndims_out = len(shape_out)
if ndims_out == 1:
tensor = ng.flatten(data)
elif ndims_out == 2:
cumprods = list(np.cumprod(shape_in))
flatten_at_idx = cumprods.index(shape_out[0]) + 1
tensor = ng.flatten_at(data, flatten_at_idx)
else:
raise NotImplementedError(
'Reshape can only support flatten to 1d or 2d.')
return ng.cast_axes(tensor, make_pos_axes(shape_out))
ax.Y.length = nout
ax.Y.name = "characters"
# Create the network
logger.debug("Creating deepspeech2 model")
ds2 = Deepspeech(args.nfilters,
args.filter_width,
args.str_w,
nbands,
args.depth,
args.hidden_size,
batch_norm=args.batch_norm)
output = ds2(inputs["audio"], spatial_axes={"H": "frequency", "W": "time"})
# set up ctc loss
loss = ng.ctc(output, ng.flatten(inputs["char_map"]),
ng.flatten(inputs["audio_length"]),
ng.flatten(inputs["char_map_length"]))
optimizer = GradientDescentMomentum(
args.lr,
momentum_coef=args.momentum,
gradient_clip_norm=args.gradient_clip_norm,
nesterov=args.nesterov)
start = time.time()
updates = optimizer(loss)
stop = time.time()
logger.debug("Optimizer graph creation took {} seconds".format(stop -
start))
mean_cost = ng.sequential([updates, ng.mean(loss, out_axes=())])
inference = True
inputs = train_set.make_placeholders()
ax.Y.length = nout
ax.Y.name = "characters"
# Create the network
logger.debug("Creating deepspeech2 model")
ds2 = Deepspeech(args.nfilters, args.filter_width, args.str_w, nbands,
args.depth, args.hidden_size, batch_norm=args.batch_norm)
output = ds2(inputs["audio"], spatial_axes={"H": "frequency", "W": "time"})
# set up ctc loss
loss = ng.ctc(output,
ng.flatten(inputs["char_map"]),
ng.flatten(inputs["audio_length"]),
ng.flatten(inputs["char_map_length"]))
optimizer = GradientDescentMomentum(args.lr,
momentum_coef=args.momentum,
gradient_clip_norm=args.gradient_clip_norm,
nesterov=args.nesterov)
start = time.time()
updates = optimizer(loss)
stop = time.time()
logger.debug("Optimizer graph creation took {} seconds".format(stop - start))
mean_cost = ng.sequential([updates, ng.mean(loss, out_axes=())])
# Create computation and initialize the transformer to allocate weights
