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是填充過后的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)
        # 分類類別數(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