def get_steps(x, time_axis, backward=False):
time_iter = list(range(time_axis.length))
if backward:
time_iter = reversed(time_iter)
if isinstance(x, dict):
return [{k: ng.slice_along_axis(x[k], time_axis, i) for k in x.keys()} for i in time_iter]
else:
return [ng.slice_along_axis(x, time_axis, i) for i in time_iter]
def run_inference(self, out_axes, init_states, **kwargs):
if self.celltype == 'LSTM':
init_states = [(state, ng.constant(0., state.axes))
for state in init_states]
one_time_axis = ng.make_axis(1, name="REC")
time_axis = out_axes.recurrent_axis()
batch_axis = out_axes.batch_axis()
feature_axis = (out_axes - [time_axis, batch_axis])[0]
outputs = [ng.constant(0., [batch_axis, one_time_axis, feature_axis])]
hidden_states = init_states
for timestep in range(time_axis.length):
in_obj = outputs[-1]
# Compute the next hidden/cell states for the recurrent layers
next_hidden_states = []
for i, l in enumerate(self.layers[:-1]):
if i < len(hidden_states):
init_state = hidden_states[i]
else:
init_state = None
if self.celltype == 'LSTM':
h, c = l(in_obj,
init_state=init_state,
return_cell_state=True)
in_obj = h
h = ng.slice_along_axis(h, one_time_axis, 0)
c = ng.slice_along_axis(c, one_time_axis, 0)
next_hidden_states.append((h, c))
else:
h = l(in_obj, init_state=init_state)
in_obj = h
h = ng.slice_along_axis(h, one_time_axis, 0)
next_hidden_states.append((h, c))
hidden_states = next_hidden_states
# Compute the output of the affine layer
in_obj = self.layers[-1](in_obj)
outputs.append(in_obj)
# Get rid of the initial 0 input
outputs = outputs[1:]
outputs = [
ng.slice_along_axis(output, one_time_axis, 0) for output in outputs
]
outputs = ng.stack(outputs, time_axis)
outputs = ng.axes_with_order(outputs, out_axes)
return outputs
def _build_module(self, input_layer):
# Dueling Network
# state value tower - V
output_axis = ng.make_axis(self.num_actions, name='q_values')
state_value = neon.Sequential([
neon.Affine(nout=256,
activation=neon.Rectlin(),
weight_init=self.weights_init,
bias_init=self.biases_init),
neon.Affine(nout=1,
weight_init=self.weights_init,
bias_init=self.biases_init)
])(input_layer)
# action advantage tower - A
action_advantage_unnormalized = neon.Sequential([
neon.Affine(nout=256,
activation=neon.Rectlin(),
weight_init=self.weights_init,
bias_init=self.biases_init),
neon.Affine(axes=output_axis,
weight_init=self.weights_init,
bias_init=self.biases_init)
])(input_layer)
action_advantage = action_advantage_unnormalized - ng.mean(
action_advantage_unnormalized)
repeated_state_value = ng.expand_dims(
ng.slice_along_axis(state_value, state_value.axes[0], 0),
output_axis, 0)
# merge to state-action value function Q
self.output = repeated_state_value + action_advantage
def _build_module(self, input_layer):
# This is almost exactly the same as Dueling Network but we predict the future measurements for each action
multistep_measurements_size = self.measurements_size[0] * self.num_predicted_steps_ahead
# actions expectation tower (expectation stream) - E
with name_scope("expectation_stream"):
expectation_stream = neon.Sequential([
neon.Affine(nout=256, activation=neon.Rectlin(),
weight_init=self.weights_init, bias_init=self.biases_init),
neon.Affine(nout=multistep_measurements_size,
weight_init=self.weights_init, bias_init=self.biases_init)
])(input_layer)
# action fine differences tower (action stream) - A
with name_scope("action_stream"):
action_stream_unnormalized = neon.Sequential([
neon.Affine(nout=256, activation=neon.Rectlin(),
weight_init=self.weights_init, bias_init=self.biases_init),
neon.Affine(nout=self.num_actions * multistep_measurements_size,
weight_init=self.weights_init, bias_init=self.biases_init),
neon.Reshape((self.num_actions, multistep_measurements_size))
])(input_layer)
action_stream = action_stream_unnormalized - ng.mean(action_stream_unnormalized)
repeated_expectation_stream = ng.slice_along_axis(expectation_stream, expectation_stream.axes[0], 0)
repeated_expectation_stream = ng.expand_dims(repeated_expectation_stream, output_axis, 0)
# merge to future measurements predictions
self.output = repeated_expectation_stream + action_stream
if args.emb is True:
# encoder input embedding
hidden_feature_axis = ng.make_axis(length=args.hs, name='hidden_feature_axis')
feature_axis = ng.make_axis(length=num_features, name='feature_axis')
W_emb = ng.variable(axes=[hidden_feature_axis, feature_axis], initial_value=init)
emb_enc_inputs = ng.dot(W_emb, inputs['inp_txt'])
# decoder input embedding
emb_dec_input = []
ax.N.length = args.batch_size
for i in range(ax.N.length):
# for each iteration, permute (by true label)
# encoder input embedding for teacher forcing input to decoder
emb_enc_input = ng.slice_along_axis(emb_enc_inputs, axis=ax.N, idx=i)
tmp_axis_1 = ng.make_axis(length=time_steps, name='tmp_axis_1')
emb_enc_input_tmp = ng.cast_axes(emb_enc_input,
ng.make_axes([hidden_feature_axis, tmp_axis_1]))
perm = ng.slice_along_axis(inputs['tgt_txt'], axis=ax.N, idx=i)
one_hot_target_tmp = ng.one_hot(perm, axis=tmp_axis_1)
emb_dec_input.append(ng.dot(emb_enc_input_tmp, one_hot_target_tmp))
emb_dec_inputs = ng.stack(emb_dec_input, axis=ax.N, pos=1)
enc_input = emb_enc_inputs
dec_input = emb_dec_inputs
else:
enc_input = inputs['inp_txt']
train_prob = seq1(inputs['inp_txt'])
train_loss = ng.cross_entropy_multi(train_prob,
ng.one_hot(inputs['tgt_txt'], axis=ax.Y),
usebits=True)
batch_cost = ng.sequential(
[optimizer(train_loss),
ng.mean(train_loss, out_axes=())])
train_outputs = dict(batch_cost=batch_cost)
with Layer.inference_mode_on():
inference_prob = seq1(inputs['inp_txt'])
errors = ng.not_equal(ng.argmax(inference_prob, reduction_axes=[ax.Y]),
inputs['tgt_txt'])
errors_last_char = ng.slice_along_axis(errors, ax.REC, time_steps - 1)
eval_loss = ng.cross_entropy_multi(inference_prob,
ng.one_hot(inputs['tgt_txt'], axis=ax.Y),
usebits=True)
eval_outputs = dict(cross_ent_loss=eval_loss,
misclass_pct=errors,
misclass_last_pct=errors_last_char)
# Now bind the computations we are interested in
with closing(ngt.make_transformer()) as transformer:
train_computation = make_bound_computation(transformer, train_outputs,
inputs)
loss_computation = make_bound_computation(transformer, eval_outputs,
inputs)
# Axis with length of batch size
N = ng.make_axis(length=params_dict['batch_size'], name='N')
# Axis with length of max question
REC = ng.make_axis(length=max_question, name='REC')
# Axis with length of hidden unit size
F = ng.make_axis(length=hidden_size, name='F')
# Axis with length of embedding size
F_embed = ng.make_axis(length=300, name='F_embed')
# Axis with length 1
dummy_axis = ng.make_axis(length=1, name='dummy_axis')
# Axis with length of answer span
span = ng.make_axis(length=2, name='span')
# Set up drop out layer
dropout_val = ng.slice_along_axis(inputs['dropout_val'], N, 0)
dropout_1 = Dropout_Modified(keep=dropout_val)
dropout_2 = Dropout_Modified(keep=dropout_val)
drop_pointer = ng.maximum(dropout_val, ng.constant(const=0.8, axes=[]))
dropout_3 = Dropout_Modified(keep=drop_pointer)
dropout_4 = Dropout_Modified(keep=drop_pointer)
# Constants required for masking
const_LSTM = ng.constant(axes=[F, dummy_axis], const=1)
const_loss = ng.constant(axes=[ax.Y, dummy_axis], const=1)
const_LSTM_embed = ng.constant(axes=[F_embed, dummy_axis], const=1)
# Create masks
reorder_para_mask = ng.axes_with_order(
inputs['para_len'], axes=[
dummy_axis, inputs['para_len'].axes[2], N])
def get_steps(x, time_axis, backward=False):
time_iter = list(range(time_axis.length))
if backward:
time_iter = reversed(time_iter)
return [ng.slice_along_axis(x, time_axis, i) for i in time_iter]
def get_steps(x, time_axis):
return [
ng.slice_along_axis(x, time_axis, i)
for i in range(time_axis.length)
]
# Make Required Axes
# Axis with length of batch size
N = ng.make_axis(length=params_dict['batch_size'], name='N')
# Axis with length of max question
REC = ng.make_axis(length=max_question, name='REC')
# Axis with length of hidden unit size
F = ng.make_axis(length=hidden_size, name='F')
# Axis with length of embedding size
F_embed = ng.make_axis(length=300, name='F_embed')
# Axis with length 1
dummy_axis = ng.make_axis(length=1, name='dummy_axis')
# Axis with length of answer span
span = ng.make_axis(length=2, name='span')
# Set up drop out layer
dropout_val = ng.slice_along_axis(inputs['dropout_val'], N, 0)
dropout_1 = Dropout_Modified(keep=dropout_val)
dropout_2 = Dropout_Modified(keep=dropout_val)
drop_pointer = ng.maximum(dropout_val, ng.constant(const=0.8, axes=[]))
dropout_3 = Dropout_Modified(keep=drop_pointer)
dropout_4 = Dropout_Modified(keep=drop_pointer)
# Constants required for masking
const_LSTM = ng.constant(axes=[F, dummy_axis], const=1)
const_loss = ng.constant(axes=[ax.Y, dummy_axis], const=1)
const_LSTM_embed = ng.constant(axes=[F_embed, dummy_axis], const=1)
# Create masks
reorder_para_mask = ng.axes_with_order(
inputs['para_len'], axes=[dummy_axis, inputs['para_len'].axes[2], N])