def forward(self, x, y, idx2tag=None):
packx = rnn.pack_sequence(x, enforce_sorted=False)
packy = rnn.pack_sequence(y, enforce_sorted=False)
embed = rnn.PackedSequence(self.embed(packx.data),
packx.batch_sizes,
sorted_indices=packx.sorted_indices,
unsorted_indices=packx.unsorted_indices)
sq_out, _ = self.lstm(embed)
sq_out = self.dropout(sq_out.data)
logits = self.classifer(sq_out)
if idx2tag is None:
loss_fct = nn.CrossEntropyLoss()
loss = loss_fct(logits, packy.data)
return loss
else:
logits.detach()
pred = torch.argmax(logits, 1)
pred_rec, _ = rnn.pad_packed_sequence(rnn.PackedSequence(
pred,
packx.batch_sizes,
sorted_indices=packx.sorted_indices,
unsorted_indices=packx.unsorted_indices),
batch_first=True)
pred_rec = pred_rec.to('cpu').numpy()
true = [[idx2tag[w.item()] for w in l] for l in y]
pred_f = [[idx2tag[w] for w in p[:len(s)]]
for s, p in zip(y, pred_rec)]
return true, pred_f
def forward(self, hidden, *packed_inputs):
"""
Applies input module to data in packed_inputs,
then applies the RNN layers,
Applies output module to output data of rnn,
Returns packed output sequence and final hidden state of RNN.
"""
batch_sizes = packed_inputs[0].batch_sizes
if self.input_module != None:
rnn_input = self.input_module(*[p.data for p in packed_inputs])
rnn_input = utils.PackedSequence(rnn_input, batch_sizes)
else:
rnn_input = packed_inputs[0]
rnn_output, hidden = self.rnn(rnn_input, hidden)
if self.output_module != None:
output = self.output_module(rnn_output.data)
output = utils.PackedSequence(output, batch_sizes)
else:
output = rnn_output
return output, hidden
def validation_step(self, batch, batch_idx): # OPTIONAL
src, tgt, shifted_tgt = batch
# print(f"batch:{batch_idx}, y:{tgt.size()}, x:{src.size()}")
output = self.forward(src, shifted_tgt)
loss = self.criterion(output.data, tgt.data) # both are packed
# loss = self.criterion(output.view(-1, output.shape[-1]), tgt.view(-1))
# metrics
preds = torch.argmax(output.data, dim=-1)
# preds = elementwise_apply(torch.argmax, output, -1)
(acc, tp, fp, fn) = metrics.acc_cm(preds, tgt.data,
self.enc.vocab_size)
preds_pad, _ = rnn.pad_packed_sequence(
rnn.PackedSequence(preds, output.batch_sizes),
batch_first=True,
padding_value=text_encoder.PAD_ID)
tgts_pad, _ = rnn.pad_packed_sequence(
tgt, batch_first=True, padding_value=text_encoder.PAD_ID)
bleu = metrics.compute_bleu(tgts_pad.tolist(), preds_pad.tolist())
return {
'val_loss': loss,
'val_acc': acc,
'tp': tp,
'fp': fp,
'fn': fn,
'bleu': bleu
}
def forward_on_packed_sequence(self,
packed_sequence,
stops_pre_filtered_flag=False):
data, *other = packed_sequence
new_data = self.forward(data, stops_pre_filtered_flag)
new_packed = rnn.PackedSequence(new_data, *other)
return new_packed
def forward(self, x, h=None):
if isinstance(x, rnn_utils.PackedSequence):
x_emb = self.embeddings(x.data)
rnn_in = rnn_utils.PackedSequence(data=x_emb, batch_sizes=x.batch_sizes)
else:
rnn_in = self.embeddings(x)
return self.rnn(rnn_in, h)
def forwards_pos(self, packed_hidden: rnn.PackedSequence) -> List[TT]:
"""Calculate the POS scores for the individual words."""
# Each element of the .data attribute of the hidden packed sequence
# should now match the input size of the linear scoring layer
assert packed_hidden.data.shape[1] == self.linear_layer.in_features
# Apply the linear layer to each element of `packed_hidden.data`
# individually
scores_data = self.linear_layer(packed_hidden.data)
# Recreate a packed sequence out of the resulting scoring data;
# this is somewhat low-level and not really recommended by PyTorch
# documentation -- do you know a better way?
packed_scores = rnn.PackedSequence(
scores_data,
batch_sizes=packed_hidden.batch_sizes,
sorted_indices=packed_hidden.sorted_indices,
unsorted_indices=packed_hidden.unsorted_indices)
# Pad the resulting packed sequence
padded_scores, padded_len = rnn.pad_packed_sequence(packed_scores,
batch_first=True)
# Convert the padded representation to a list of score tensors
scores = []
for sco, n in zip(padded_scores, padded_len):
scores.append(sco[:n])
# Finally, return the scores
return scores
def prepend_tensor_to_start_of_packed_seq(packed_seq: rnn.PackedSequence,
value_to_add):
"""
This function shifts the whole sequence down and adds value_to_add to the start.
"""
data, batch_sizes, *others = packed_seq
# We're gonna be a bit cheeky and construct a Packed Sequence manually at the bottom of this function -- which the
# docs tell us not to do but have seen others do it, eg
# https://github.com/pytorch/pytorch/issues/8921#issuecomment-400552029
# Originally we coded this in PyTorch 1.0 and PackedSequence was a thinner wrapper on a NamedTuple
# so continue to check that we are still using enforce_sorted=True Packed Sequences
if len(others):
assert others[0] is None
assert others[1] is None
num_in_first_batch = batch_sizes[0]
front = torch.zeros_like(data[:num_in_first_batch])
front[...] = value_to_add
new_packed_seq_data = torch.cat([front, data], dim=0)
new_length_at_beginning = batch_sizes[:1].clone()
new_packed_seq = rnn.PackedSequence(
new_packed_seq_data,
torch.cat([new_length_at_beginning, packed_seq.batch_sizes]))
return new_packed_seq
def alpha_packed(self, scores: rnn.PackedSequence) -> rnn.PackedSequence:
"""Variant of `alpha` which works on packed sequences."""
# Additional check
class_num = self.class_num()
assert class_num == scores.data.shape[1]
# Tensor for results, accounts for the input scores
alpha = scores.data.clone()
# Transposed transition weights matrix
T_t = self.T.t()
# Sizes and start positions of the batches
size = scores.batch_sizes
start = batch_start(size)
# Loop over the batches (except the first), left to right
for i in range(1, len(size)):
# Alpha corresponding to the previous batch
alpha_prev = alpha[start[i - 1]:start[i - 1] + size[i]]
# For each row alpha_row in alpha_prev, create a separate instance
# of T_t and add alpha_row to each row in this instance of T_t.
# As a result, we should obtain B copies of T_t, where B is
# the size of the batch.
sum_tensor = add_to_each(alpha_prev, T_t)
# Apply logsumexp in a batch
alpha[start[i]:start[i]+size[i]] += \
torch.logsumexp(sum_tensor, dim=2)
# Return the result
return rnn.PackedSequence(alpha, scores.batch_sizes,
scores.sorted_indices,
scores.unsorted_indices)
def concat_packed_seq_features(packed_seqs):
data = []
batch_sizes = None
for p_seq in packed_seqs:
d, b_sizes, *others = p_seq
# We're gonna be a bit cheeky and construct a Packed Sequence manually at the bottom of this function -- which the
# docs tell us not to do but have seen others do it, eg
# https://github.com/pytorch/pytorch/issues/8921#issuecomment-400552029
# Originally we coded this in PyTorch 1.0 and PackedSequence was a thinner wrapper on a NamedTuple
# so continue to check that we are still using enforce_sorted=True Packed Sequences
if len(others):
assert others[0] is None
assert others[1] is None
if batch_sizes is None:
batch_sizes = b_sizes
else:
if not torch.all(batch_sizes == b_sizes):
raise RuntimeError(
"Trying to concat incompatible packed sequences")
data.append(d)
data = torch.cat(data, dim=-1)
return rnn.PackedSequence(data, batch_sizes)
def forward(self, output, context):
# The idiomatic way to perform an explicit typecheck in Python is to use isinstance(x, Y) rather than type(x) == Y, type(x) is Y.
# Though there are unusual situations where these give different results.
# To check whether the output is type of PackedSequence.
if(isinstance(output, rnn_utils.PackedSequence)):
unpack_output, _ = rnn_utils.pad_packed_sequence(output, batch_first=True)
else:
unpack_output = output
batch_size = unpack_output.size(0)
hidden_size = unpack_output.size(2)
context_trans = context.view(1, -1, context.size(1))
input_size = context_trans.size(1)
# for idx1, temp1 in enumerate(unpack_output[0]):
# for idx, temp in enumerate(context_trans[0]):
# print('o'+str(idx1)+',c'+str(idx)+': '+str(torch.dot(temp1, temp)))
# (batch, out_len, dim) * (batch, in_len, dim) -> (batch, out_len, in_len)
attn = torch.bmm(unpack_output, context_trans.transpose(1, 2))
if self.mask is not None:
attn.data.masked_fill_(self.mask, -float('inf'))
attn = F.softmax(attn.view(-1, input_size), dim=1).view(batch_size, -1, input_size)
# (batch, out_len, in_len) * (batch, in_len, dim) -> (batch, out_len, dim)
mix = torch.bmm(attn, context_trans)
# concat -> (batch, out_len, 2*dim)
combined = torch.cat((mix, unpack_output), dim=2)
# output -> (batch, out_len, dim)
output_result = torch.tanh(self.linear_out(combined.view(-1, 2 * hidden_size))).view(-1, hidden_size)
# Transform result into PackedSequence format.
packed_output_result = rnn_utils.PackedSequence(output_result, output.batch_sizes.detach()) if isinstance(output, rnn_utils.PackedSequence) else output_result.view(1, -1, hidden_size)
return packed_output_result, attn
def nlog_like_of_obs(self, obs: rnn.PackedSequence):
"""
Here we calculate the negative log likelihood of the sequence. For each we feed in the previous observation
ie if you use this function during training then doing teacher forcing.
"""
# Set up the ground truth inputs from previous time-steps to be fed into the bottom of the RNN
symbol_seq_packed_minus_last = torch_utils.remove_last_from_packed_seq(
obs)
embeddings = self.embedder.forward_on_packed_sequence(
symbol_seq_packed_minus_last, stops_pre_filtered_flag=True)
inputs = torch_utils.prepend_tensor_to_start_of_packed_seq(
embeddings, mchef_config.SOS_TOKEN)
# Feed the emebeddings through the network
initial_hidden = self._initial_hidden_after_update
outputs, _ = self.gru(inputs, initial_hidden)
outputs_mapped = self.mlp_out(outputs.data)
self.decoder_top.update(outputs_mapped)
# Now work out the nll for each element of each sequence and then sum over the whole sequence length.
nll_per_obs = self.decoder_top.nlog_like_of_obs(obs.data)
nll_packed = rnn.PackedSequence(nll_per_obs, *obs[1:])
nll_padded, _ = rnn.pad_packed_sequence(nll_packed,
batch_first=True,
padding_value=0.0)
nll_per_seq = nll_padded.sum(
dim=tuple(range(1, len(nll_padded.shape))))
return nll_per_seq
def forward(self, gru_x, att_score, seq_len):
'''
:param gru_x: [B, max_seq_len, input_dim]
:param att_score: [B, max_seq_len]
:param seq_len: [B]
:return:
'''
#define
dtype = gru_x.dtype
device = gru_x.device
#data [sum(seq_len), input_dim]
#batch_sizes sorted(seq_len)
pack_data, batch_sizes, sorted_indices, unsorted_indices = rnn_utils.pack_padded_sequence(
gru_x, seq_len.cpu(), batch_first=True, enforce_sorted=False)
#sorce [sum(seq_len)]
pack_score, _, _, _ = rnn_utils.pack_padded_sequence(
att_score, seq_len.cpu(), batch_first=True,
enforce_sorted=False) #out [B, max_seq_len, hidden_dim]
#gru init
h0 = torch.zeros(batch_sizes[0],
self.hidden_size,
dtype=dtype,
device=device)
#output
#pack_out [sum(seq_len), hidden_dim]
data_out = torch.zeros(pack_data.size(0),
self.hidden_size,
dtype=dtype,
device=device)
#times
begin = 0
ht = h0
for batch in batch_sizes:
i_x = pack_data[begin:begin + batch] #[batch, input_dim]
s_x = pack_score[begin:begin + batch] #[batch]
h_x = ht[:batch] #[batch, hidden_dim]
h_t_plus_1 = self.gru(i_x, s_x, h_x) #[batch, hidden_dim]
data_out[begin:begin + batch] = h_t_plus_1
ht = h_t_plus_1
begin += batch
pack_out = rnn_utils.PackedSequence(data_out, batch_sizes,
sorted_indices, unsorted_indices)
out, _ = rnn_utils.pad_packed_sequence(
pack_out,
batch_first=True,
padding_value=0.0,
total_length=gru_x.size(1)) #out [B, max_seq_len, hidden_dim]
return out
def forward(self, x, lengths):
px = self.traced_packer(x, lengths)
o = self.lstm(rnn_utils.PackedSequence(*px))
# print(o)
px = self.traced_unpacker(o[0].data, o[0].batch_sizes)
# Assume only one event
y1 = self.pred(px[0])
y = self.out(y1)
return y
def forward(self, G_t, ret_hid=False, *args):
h, hidden = self.theta_net(G_t, *args)
discrete_features = rnn.PackedSequence(
self.discrete_feature_net(h.data), h.batch_sizes)
adjacencies = rnn.PackedSequence(self.adjacency_feature_net(h.data),
h.batch_sizes)
if hasattr(self, 'continuous_feature_net'):
continuous_features = rnn.PackedSequence(
self.continuous_feature_net(h.data), h.batch_sizes)
if ret_hid:
return discrete_features, continuous_features, adjacencies, hidden
return discrete_features, continuous_features, adjacencies
else:
if ret_hid:
return discrete_features, adjacencies, hidden
return discrete_features, adjacencies
def score_packed(self, seq: PackedSequence) -> PackedSequence:
# Apply the linear layer to each hidden vector in the packed sequence
# in order the get the final scores
scores_data = self.score_layer(seq.data)
# Create a packed sequence from the result. According to PyTorch
# documentation, this should be never done, but how to do that
# differently? Of course we *could* do that after padding the
# packed sequence, but this seems suboptimal.
return rnn.PackedSequence(scores_data, seq.batch_sizes,
seq.sorted_indices, seq.unsorted_indices)
def pack_batch_no_out(self, batch, embeddings, device="cpu"):
# Asserting statements is a convenient way to insert debugging assertions into a program.
# To guarantee that the batch is a list.
assert isinstance(batch, list)
# The format of batch is a list of tuple: ((tuple),[[list of token ID list]])
# A lambda function is a small anonymous function, the example is as following.
# x = lambda a, b: a * b
# print(x(5, 6))
batch.sort(key=lambda s: len(s[0]), reverse=True)
# input_idx：Each row is corresponding to one input sentence.
# output_idx：ach row is corresponding to a list of several output sentences.
# zip wants a bunch of arguments to zip together, but what you have is a single argument (a list, whose elements are also lists).
# The * in a function call "unpacks" a list (or other iterable), making each of its elements a separate argument.
# For list p = [[1,2,3],[4,5,6]];
# So without the *, you're doing zip( [[1,2,3],[4,5,6]] ). With the *, you're doing zip([1,2,3], [4,5,6]) = [(1, 4), (2, 5), (3, 6)].
input_idx, output_idx = zip(*batch)
# create padded matrix of inputs
# map() function returns a list of the results after applying the given function to each item of a given iterable (list, tuple etc.)
# For example:
# numbers = (1, 2, 3, 4)
# result = map(lambda x: x + x, numbers)
# print(list(result))
# Output: {2, 4, 6, 8}
lens = list(map(len, input_idx))
input_mat = np.zeros((len(batch), lens[0]), dtype=np.int64)
for idx, x in enumerate(input_idx):
input_mat[idx, :len(x)] = x
input_v = torch.tensor(input_mat).to(device)
input_v = input_v.cuda()
# The padded sequence is the transposed matrix which is ``B x T x *``,
# where `T` is the length of the longest sequence and `B` is the batch size.
# Following the matrix is the list of lengths of each sequence in the batch (also in transposed format).
# For instance:
# [ a b c c d d d ]
# [ a b c d ]
# [ a b c ]
# [ a b ]
# could be transformed into [a,a,a,a,b,b,b,b,c,c,c,c,d,d,d,d] with batch size [4,4,3,2,1,1,1].
input_seq = rnn_utils.pack_padded_sequence(input_v,
lens,
batch_first=True)
input_seq = input_seq.cuda()
r = embeddings(input_seq.data)
# r: the [B x T x dimension] matrix of the embeddings of the occurred words in input sequence.
# The order is followed by the order in input_seq.
# Which is transforming [a,a,a,a,b,b,b,b,c,c,c,c,d,d,d,d] into [embedding(a), embedding(a), ..., embedding(d), embedding(d)]
r = r.cuda()
# For instance, given data ``abc`` and `x`
# the :class:`PackedSequence` would contain data ``axbc`` with ``batch_sizes=[2,1,1]``.
# emb_input_seq is [B x T x dimension] matrix of the embeddings of the occurred words in input sequence with the batch size.
# For instance, emb_input_seq is the padded data: [embedding(a), embedding(a), ..., embedding(d), embedding(d)] with batch size [4,4,3,2,1,1,1].
emb_input_seq = rnn_utils.PackedSequence(r, input_seq.batch_sizes)
emb_input_seq = emb_input_seq.cuda()
return emb_input_seq, input_idx, output_idx
def forward(self, x, h_0=None):
is_packed = isinstance(x, _rnn_utils.PackedSequence)
if is_packed:
y = self.ff(x.data)
y = _rnn_utils.PackedSequence(y,
batch_sizes=x.batch_sizes,
sorted_indices=x.sorted_indices,
unsorted_indices=x.unsorted_indices)
else:
y = self.ff(x)
return y, _torch.Tensor() #no hidden state.
def forward(self, input, lengths=None):
if lengths is None:
N, T = input.size(0), input.size(1)
lengths = [T for _ in range(N)]
output = self.CBHG(input, lengths).contiguous().view(N * T, -1)
output = self.projection(output).view(N, T, -1)
else:
output = self.CBHG(input, lengths)
output = rnn.pack_padded_sequence(output, lengths, True)
output = rnn.PackedSequence(self.projection(output.data),
output.batch_sizes)
output, _ = rnn.pad_packed_sequence(output, True)
return output
def forward(self, x):
x_img, x_text = x
assert isinstance(x_text, rnn_utils.PackedSequence)
# deal with text data
emb_out = self.emb(x_text.data)
emb_out_seq = rnn_utils.PackedSequence(data=emb_out, batch_sizes=x_text.batch_sizes)
rnn_out, rnn_h = self.rnn(emb_out_seq)
# extract image features
fx = x_img.float() / 256
conv_out = self.conv(fx).view(fx.size()[0], -1)
feats = self._concat_features(conv_out, rnn_h)
return self.policy(feats), self.value(feats)
def pack_batch_no_out(batch, embeddings, device="cpu"):
assert isinstance(batch, list)
# Sort descending (CuDNN requirements)
batch.sort(key=lambda s: len(s[0]), reverse=True)
input_idx, output_idx = zip(*batch)
# create padded matrix of inputs
lens = list(map(len, input_idx))
input_mat = np.zeros((len(batch), lens[0]), dtype=np.int64)
for idx, x in enumerate(input_idx):
input_mat[idx, :len(x)] = x
input_v = torch.tensor(input_mat).to(device)
input_seq = rnn_utils.pack_padded_sequence(input_v, lens, batch_first=True)
# lookup embeddings
r = embeddings(input_seq.data)
emb_input_seq = rnn_utils.PackedSequence(r, input_seq.batch_sizes)
return emb_input_seq, input_idx, output_idx
def forward(self, x):
# use pack_sequence for the best performance
x = rnn_utils.pack_sequence(x)
#x = rnn_utils.pad_sequence(x).cuda()
embedded = rnn_utils.PackedSequence(self.embedding(x.data),
x.batch_sizes)
#embedded = self.embedding(x)
_, hidden = self.rnn(embedded)
#hidden, _ = rnn_utils.pad_packed_sequence(hidden)
hidden = self.dropout(torch.cat((hidden[-2], hidden[-1]), dim=1))
out = self.fc1(hidden)
return out.squeeze()
def pack_batch_no_out(batch, embeddings, device="cpu"):
# Assert statements are a convenient way to insert debugging assertions into a program.
# To guarantee that the batch is a list.
assert isinstance(batch, list)
# The format of batch is a list of tuple: ((tuple),[[list of token ID list]])
# A lambda function is a small anonymous function, the example is as following.
# x = lambda a, b: a * b
# print(x(5, 6))
# Sort descending (CuDNN requirements) batch中第一个元素为最长的句子；
batch.sort(key=lambda s: len(s[0]), reverse=True)
# input_idx：一个batch的输入句子的tokens对应的ID矩阵；Each row is corresponding to one input sentence.
# output_idx：一个batch的输出句子的tokens对应的ID矩阵；Each row is corresponding to a list of several output sentences.
input_idx, output_idx = zip(*batch)
# create padded matrix of inputs
# map() function returns a list of the results after applying the given function to each item of a given iterable (list, tuple etc.)
# For example:
# numbers = (1, 2, 3, 4)
# result = map(lambda x: x + x, numbers)
# print(list(result))
# Output: {2, 4, 6, 8}
# 建立长度词典，为batch中每一个元素的长度；
lens = list(map(len, input_idx))
# 以最长的句子来建立batch*最长句子长度的全0矩阵；
input_mat = np.zeros((len(batch), lens[0]), dtype=np.int64)
# 以固定的句子长度来建立batch*固定长句子长度的全0矩阵；
# input_mat = np.zeros((len(batch), PADDING_SIZE), dtype=np.int64)
# 将batch中每个句子的tokens对应的ID向量填入全0矩阵完成padding；
# idx：index，x：token ID 组成的向量；
for idx, x in enumerate(input_idx):
input_mat[idx, :len(x)] = x
# 将padding后的矩阵转换为tensor matrix；
input_v = torch.tensor(input_mat).to(device)
input_v = input_v.cuda()
# 封装成PackedSequence类型的对象；
input_seq = rnn_utils.pack_padded_sequence(input_v, lens, batch_first=True)
input_seq = input_seq.cuda()
# lookup embeddings；embeddings为模型已经建立的词向量矩阵；
r = embeddings(input_seq.data)
r = r.cuda()
# 加入了词嵌入的input_seq；
emb_input_seq = rnn_utils.PackedSequence(r, input_seq.batch_sizes)
emb_input_seq = emb_input_seq.cuda()
return emb_input_seq, input_idx, output_idx
def forward_(self, x):
# x: N, inc, inn
# import pdb; pdb.set_trace()
N = x.shape[0]
embedding_weight = x.view(-1, self.inn).t() # inn, (N * inc)
embedding_output = F.embedding(self.A_packed_data, embedding_weight) # L, (N * inc)
masked_embedding_output, _ = rnn.pad_packed_sequence(
rnn.PackedSequence(
data=self.mask_weight_packed * embedding_output if use_mask_weight else embedding_output,
batch_sizes=self.A_packed_batch_sizes,
),
batch_first=True,
) # outn, D, (N * inc)
pooled_masked_embedding_output = masked_embedding_output.sum(dim=1) # outn, (N * inc) but in wrong order.
gathered_pooled_masked_embedding_output = pooled_masked_embedding_output[self.original_idx] # outn, (N * inc) in correct order.
scaled_gpmeo = gathered_pooled_masked_embedding_output * self.weight.expand(self.outn, N, self.inc).contiguous().view(self.outn, N * self.inc)
pre_channel_transform = scaled_gpmeo.view(-1, self.inc) # .view(self.outn, N, self.inc).contiguous() # Flattened for linear
almost_y = self.channel_transform(pre_channel_transform).view(self.outn, N, self.outc) # outn, N, outc
y = almost_y.contiguous().permute(1, 2, 0).contiguous()
return y + self.bias
def _forward_packed(self, some_data):
lengths = None
x = None
y = None
x = some_data.data
lengths = some_data.batch_sizes
if x.ndim == 2 and x.shape[-1] == 1:
x = _torch.flatten(x)
assert x.ndim == 1, "Packed data should only have 1 dimension before embedding"
y = self.ff(x)
return _rnn_utils.PackedSequence(
y,
batch_sizes=some_data.batch_sizes,
sorted_indices=some_data.sorted_indices,
unsorted_indices=some_data.unsorted_indices)
def forward(self,
input_enc,
input_attW_enc,
input_dec,
lengths_enc,
hidden_att=None,
hidden_dec1=None,
hidden_dec2=None):
N = input_dec.size(0)
out_att = self.prenet(input_dec).unsqueeze(1) # N x O_dec -> N x 1 x H
out_att, hidden_att = self.gru_att(out_att, hidden_att) # N x 1 x 2H
in_attW_dec = self.linear_dec(
out_att.squeeze(1)).unsqueeze(1).expand_as(input_enc)
in_attW_dec = rnn.pack_padded_sequence(in_attW_dec, lengths_enc,
True) # N*T_enc x 2H
self.attn_weights = torch.add(
input_attW_enc, in_attW_dec.data).tanh() # N x T_enc x 2H
self.attn_weights = self.attn(self.attn_weights).exp() # N*T_enc x 1
self.attn_weights = rnn.PackedSequence(self.attn_weights,
in_attW_dec.batch_sizes)
self.attn_weights, _ = rnn.pad_packed_sequence(self.attn_weights, True)
self.attn_weights = F.normalize(self.attn_weights, 1,
1) # N x T_enc x 1
attn_applied = torch.bmm(self.attn_weights.transpose(1, 2),
input_enc) # N x 1 x 2H
out_dec = torch.cat((attn_applied, out_att), 2) # N x 1 x 4H
residual = self.short_cut(out_dec.squeeze(1)).unsqueeze(
1) # N x 1 x 2H
out_dec, hidden_dec1 = self.gru_dec1(out_dec, hidden_dec1)
residual = residual + out_dec
out_dec, hidden_dec2 = self.gru_dec2(residual, hidden_dec2)
residual = residual + out_dec
output = self.out(residual.squeeze(1)).view(N, self.r_factor, -1)
return output, hidden_att, hidden_dec1, hidden_dec2
def reduction_(self, x):
# x: N, c, inn
# output: N, c, outn
N, c, n = x.shape
embw = x.view(-1, n).t()
embo = F.embedding(self.A_packed_data, embw) # L, (N * c)
masked_embo, _ = rnn.pad_packed_sequence(
rnn.PackedSequence(
data=self.mask_weight_packed * embo,
batch_sizes=self.A_packed_batch_sizes,
),
batch_first=True,
) # outn, D, (N * c)
masked_embo = masked_embo[original_idx]
if self.reduction == "max":
pooled_masked_embo = masked_embo.max(dim=1)[0]
else:
pooled_masked_embo = masked_embo.sum(dim=1)
if self.reduction == "mean":
pooled_masked_embo = pooled_masked_embo / self.lengths_f32
return pooled_masked_embo.t().view(N, c, -1).contiguous()
def beta_packed(self, scores: rnn.PackedSequence) -> rnn.PackedSequence:
"""Variant of `beta` which works on packed sequences."""
# Additional check
class_num = self.class_num()
assert class_num == scores.data.shape[1]
# Tensor for results, accounts for the input scores
beta = scores.data.clone()
# Sizes and start positions of the batches
size = scores.batch_sizes
start = batch_start(size)
# Loop over the batches (except the last), right to left
for i in range(len(size) - 2, -1, -1):
# Beta corresponding to the next batch
beta_next = beta[start[i + 1]:start[i + 1] + size[i + 1]]
# See `add_to_each` and `alpha_packed`
sum_tensor = add_to_each(beta_next, self.T)
# Apply logsumexp in a batch
beta[start[i]:start[i]+size[i+1]] += \
torch.logsumexp(sum_tensor, dim=2)
# Return the result
return rnn.PackedSequence(beta, scores.batch_sizes,
scores.sorted_indices,
scores.unsorted_indices)
def forward(self, seq: PackedSequence) -> PackedSequence:
seq_data = self.dropout(seq.data)
return rnn.PackedSequence(seq_data, seq.batch_sizes,
seq.sorted_indices, seq.unsorted_indices)
def unpack(self, x, l):
p = rnn_utils.PackedSequence(x, l)
return rnn_utils.pad_packed_sequence(p, batch_first=True)
def apply2packed(f, ps):
return rnn_utils.PackedSequence(f(ps.data), ps.batch_sizes)
