def pad_tensor(vec, pad):
 """
 args:
  vec - tensor to pad
  pad - the size to pad to
 return:
  a new tensor padded to 'pad'
 """
 return torch.cat([vec, torch.zeros(pad - len(vec), dtype=torch.float)], dim=0).data.numpy()
class Collate:
 """
 a variant of callate_fn that pads according to the longest sequence in
 a batch of sequences
 """
 def __init__(self):
  pass
 def _collate(self, batch):
  """
  args:
batch - list of (tensor, label)
  reutrn:
xs - a tensor of all examples in 'batch' before padding like:
 '''
 [tensor([1,2,3,4]),
  tensor([1,2]),
  tensor([1,2,3,4,5])]
 '''
ys - a LongTensor of all labels in batch like:
 '''
 [1,0,1]
 '''
  """
  xs = [torch.FloatTensor(v[0]) for v in batch]
  ys = torch.LongTensor([v[1] for v in batch])
  # 獲得每個(gè)樣本的序列長(zhǎng)度
  seq_lengths = torch.LongTensor([v for v in map(len, xs)])
  max_len = max([len(v) for v in xs])
  # 每個(gè)樣本都padding到當(dāng)前batch的最大長(zhǎng)度
  xs = torch.FloatTensor([pad_tensor(v, max_len) for v in xs])
  # 把xs和ys按照序列長(zhǎng)度從大到小排序
  seq_lengths, perm_idx = seq_lengths.sort(0, descending=True)
  xs = xs[perm_idx]
  ys = ys[perm_idx]
  return xs, seq_lengths, ys
 def __call__(self, batch):
  return self._collate(batch)

train_data = Data.DataLoader(dataset=train_dataset, batch_size=32, num_workers=0, collate_fn=Collate())

from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
# x是填充過(guò)后的batch數(shù)據(jù)，seq_lengths是每個(gè)樣本的序列長(zhǎng)度
packed_input = pack_padded_sequence(x, seq_lengths, batch_first=True)

class Model(nn.Module):
 def __init__(self, in_size, hid_size, n_layer, drop=0.1, bi=False):
  super(Model, self).__init__()
  self.lstm = nn.LSTM(input_size=in_size,
hidden_size=hid_size,
num_layers=n_layer,
batch_first=True,
dropout=drop,
bidirectional=bi)
  # 分類(lèi)類(lèi)別數(shù)目為2
  self.fc = nn.Linear(in_features=hid_size, out_features=2)
 def forward(self, x):
  '''
  :param x: 變長(zhǎng)序列時(shí)，x是一個(gè)PackedSequence對(duì)象
  :return: PackedSequence對(duì)象
  '''
  # lstm_out: tensor of shape (batch, seq_len, num_directions * hidden_size)
  lstm_out, _ = self.lstm(x)  
  
  return lstm_out
model = Model()
lstm_out = model(packed_input)

# 此處lstm_out是一個(gè)PackedSequence對(duì)象
output, _ = pad_packed_sequence(lstm_out)

import torch
from torch import nn
class ImprovedRnn(nn.Module):
 def __init__(self, module, *args, **kwargs):
  assert module in (nn.RNN, nn.LSTM, nn.GRU)
  super().__init__()
  self.module = module(*args, **kwargs)
 def forward(self, input, lengths):  # input shape(batch_size, seq_len, input_size)
  if not hasattr(self, '_flattened'):
self.module.flatten_parameters()
setattr(self, '_flattened', True)
  max_len = input.shape[1]
  # enforce_sorted=False則自動(dòng)按lengths排序，并且返回值package.unsorted_indices可用于恢復(fù)原順序
  package = nn.utils.rnn.pack_padded_sequence(input, lengths.cpu(), batch_first=self.module.batch_first, enforce_sorted=False)
  result, hidden = self.module(package)
  # total_length參數(shù)一般不需要，因?yàn)閘engths列表中一般含最大值。但分布式訓(xùn)練時(shí)是將一個(gè)batch切分了，故一定要有！
  result, lens = nn.utils.rnn.pad_packed_sequence(result, batch_first=self.module.batch_first, total_length=max_len)
  return result[package.unsorted_indices], hidden  # output shape(batch_size, seq_len, rnn_hidden_size)

class TestNet(nn.Module):
 def __init__(self, word_emb, gru_in, gru_out):
  super().__init__()
  self.encode = nn.Embedding.from_pretrained(torch.Tensor(word_emb))
  self.rnn = ImprovedRnn(nn.RNN, input_size=gru_in, hidden_size=gru_out,
		  				batch_first=True, bidirectional=True)
 def forward(self, seq1, seq1_lengths, seq2, seq2_lengths):
  seq1_emb = self.encode(seq1)
  seq2_emb = self.encode(seq2)
  rnn1, hn = self.rnn(seq1_emb, seq1_lengths)
  rnn2, hn = self.rnn(seq2_emb, seq2_lengths)
  """
  此處略去rnn1和rnn2的后續(xù)計(jì)算，當(dāng)前網(wǎng)絡(luò)最后計(jì)算結(jié)果記為prediction
  """
  return prediction