import torch.nn as nn import torch from fastNLP.core.utils import seq_len_to_mask from utils import better_init_rnn import numpy as np class WordLSTMCell_yangjie(nn.Module): """A basic LSTM cell.""" def __init__(self, input_size, hidden_size, use_bias=True,debug=False, left2right=True): """ Most parts are copied from torch.nn.LSTMCell. """ super().__init__() self.left2right = left2right self.debug = debug self.input_size = input_size self.hidden_size = hidden_size self.use_bias = use_bias self.weight_ih = nn.Parameter( torch.FloatTensor(input_size, 3 * hidden_size)) self.weight_hh = nn.Parameter( torch.FloatTensor(hidden_size, 3 * hidden_size)) if use_bias: self.bias = nn.Parameter(torch.FloatTensor(3 * hidden_size)) else: self.register_parameter('bias', None) self.reset_parameters() def reset_parameters(self): """ Initialize parameters following the way proposed in the paper. """ nn.init.orthogonal(self.weight_ih.data) weight_hh_data = torch.eye(self.hidden_size) weight_hh_data = weight_hh_data.repeat(1, 3) with torch.no_grad(): self.weight_hh.set_(weight_hh_data) # The bias is just set to zero vectors. if self.use_bias: nn.init.constant(self.bias.data, val=0) def forward(self, input_, hx): """ Args: input_: A (batch, input_size) tensor containing input features. hx: A tuple (h_0, c_0), which contains the initial hidden and cell state, where the size of both states is (batch, hidden_size). Returns: h_1, c_1: Tensors containing the next hidden and cell state. """ h_0, c_0 = hx batch_size = h_0.size(0) bias_batch = (self.bias.unsqueeze(0).expand(batch_size, *self.bias.size())) wh_b = torch.addmm(bias_batch, h_0, self.weight_hh) wi = torch.mm(input_, self.weight_ih) f, i, g = torch.split(wh_b + wi, split_size_or_sections=self.hidden_size, dim=1) c_1 = torch.sigmoid(f)*c_0 + torch.sigmoid(i)*torch.tanh(g) return c_1 def __repr__(self): s = '{name}({input_size}, {hidden_size})' return s.format(name=self.__class__.__name__, **self.__dict__) class MultiInputLSTMCell_V0(nn.Module): def __init__(self, char_input_size, hidden_size, use_bias=True,debug=False): super().__init__() self.char_input_size = char_input_size self.hidden_size = hidden_size self.use_bias = use_bias self.weight_ih = nn.Parameter( torch.FloatTensor(char_input_size, 3 * hidden_size) ) self.weight_hh = nn.Parameter( torch.FloatTensor(hidden_size, 3 * hidden_size) ) self.alpha_weight_ih = nn.Parameter( torch.FloatTensor(char_input_size, hidden_size) ) self.alpha_weight_hh = nn.Parameter( torch.FloatTensor(hidden_size, hidden_size) ) if self.use_bias: self.bias = nn.Parameter(torch.FloatTensor(3 * hidden_size)) self.alpha_bias = nn.Parameter(torch.FloatTensor(hidden_size)) else: self.register_parameter('bias', None) self.register_parameter('alpha_bias', None) self.debug = debug self.reset_parameters() def reset_parameters(self): """ Initialize parameters following the way proposed in the paper. """ nn.init.orthogonal(self.weight_ih.data) nn.init.orthogonal(self.alpha_weight_ih.data) weight_hh_data = torch.eye(self.hidden_size) weight_hh_data = weight_hh_data.repeat(1, 3) with torch.no_grad(): self.weight_hh.set_(weight_hh_data) alpha_weight_hh_data = torch.eye(self.hidden_size) alpha_weight_hh_data = alpha_weight_hh_data.repeat(1, 1) with torch.no_grad(): self.alpha_weight_hh.set_(alpha_weight_hh_data) # The bias is just set to zero vectors. if self.use_bias: nn.init.constant_(self.bias.data, val=0) nn.init.constant_(self.alpha_bias.data, val=0) def forward(self, inp, skip_c, skip_count, hx): ''' :param inp: chars B * hidden :param skip_c: 由跳边得到的c, B * X * hidden :param skip_count: 这个batch中每个example中当前位置的跳边的数量，用于mask :param hx: :return: ''' max_skip_count = torch.max(skip_count).item() if True: h_0, c_0 = hx batch_size = h_0.size(0) bias_batch = (self.bias.unsqueeze(0).expand(batch_size, *self.bias.size())) wi = torch.matmul(inp, self.weight_ih) wh = torch.matmul(h_0, self.weight_hh) i, o, g = torch.split(wh + wi + bias_batch, split_size_or_sections=self.hidden_size, dim=1) i = torch.sigmoid(i).unsqueeze(1) o = torch.sigmoid(o).unsqueeze(1) g = torch.tanh(g).unsqueeze(1) alpha_wi = torch.matmul(inp, self.alpha_weight_ih) alpha_wi.unsqueeze_(1) # alpha_wi = alpha_wi.expand(1,skip_count,self.hidden_size) alpha_wh = torch.matmul(skip_c, self.alpha_weight_hh) alpha_bias_batch = self.alpha_bias.unsqueeze(0) alpha = torch.sigmoid(alpha_wi + alpha_wh + alpha_bias_batch) skip_mask = seq_len_to_mask(skip_count,max_len=skip_c.size()[1]) skip_mask = 1 - skip_mask skip_mask = skip_mask.unsqueeze(-1).expand(*skip_mask.size(), self.hidden_size) skip_mask = (skip_mask).float()*1e20 alpha = alpha - skip_mask alpha = torch.exp(torch.cat([i, alpha], dim=1)) alpha_sum = torch.sum(alpha, dim=1, keepdim=True) alpha = torch.div(alpha, alpha_sum) merge_i_c = torch.cat([g, skip_c], dim=1) c_1 = merge_i_c * alpha c_1 = c_1.sum(1, keepdim=True) # h_1 = o * c_1 h_1 = o * torch.tanh(c_1) return h_1.squeeze(1), c_1.squeeze(1) else: h_0, c_0 = hx batch_size = h_0.size(0) bias_batch = (self.bias.unsqueeze(0).expand(batch_size, *self.bias.size())) wi = torch.matmul(inp, self.weight_ih) wh = torch.matmul(h_0, self.weight_hh) i, o, g = torch.split(wh + wi + bias_batch, split_size_or_sections=self.hidden_size, dim=1) i = torch.sigmoid(i).unsqueeze(1) o = torch.sigmoid(o).unsqueeze(1) g = torch.tanh(g).unsqueeze(1) c_1 = g h_1 = o * c_1 return h_1,c_1 class MultiInputLSTMCell_V1(nn.Module): def __init__(self, char_input_size, hidden_size, use_bias=True,debug=False): super().__init__() self.char_input_size = char_input_size self.hidden_size = hidden_size self.use_bias = use_bias self.weight_ih = nn.Parameter( torch.FloatTensor(char_input_size, 3 * hidden_size) ) self.weight_hh = nn.Parameter( torch.FloatTensor(hidden_size, 3 * hidden_size) ) self.alpha_weight_ih = nn.Parameter( torch.FloatTensor(char_input_size, hidden_size) ) self.alpha_weight_hh = nn.Parameter( torch.FloatTensor(hidden_size, hidden_size) ) if self.use_bias: self.bias = nn.Parameter(torch.FloatTensor(3 * hidden_size)) self.alpha_bias = nn.Parameter(torch.FloatTensor(hidden_size)) else: self.register_parameter('bias', None) self.register_parameter('alpha_bias', None) self.debug = debug self.reset_parameters() def reset_parameters(self): """ Initialize parameters following the way proposed in the paper. """ nn.init.orthogonal(self.weight_ih.data) nn.init.orthogonal(self.alpha_weight_ih.data) weight_hh_data = torch.eye(self.hidden_size) weight_hh_data = weight_hh_data.repeat(1, 3) with torch.no_grad(): self.weight_hh.set_(weight_hh_data) alpha_weight_hh_data = torch.eye(self.hidden_size) alpha_weight_hh_data = alpha_weight_hh_data.repeat(1, 1) with torch.no_grad(): self.alpha_weight_hh.set_(alpha_weight_hh_data) # The bias is just set to zero vectors. if self.use_bias: nn.init.constant_(self.bias.data, val=0) nn.init.constant_(self.alpha_bias.data, val=0) def forward(self, inp, skip_c, skip_count, hx): ''' :param inp: chars B * hidden :param skip_c: 由跳边得到的c, B * X * hidden :param skip_count: 这个batch中每个example中当前位置的跳边的数量，用于mask :param hx: :return: ''' max_skip_count = torch.max(skip_count).item() if True: h_0, c_0 = hx batch_size = h_0.size(0) bias_batch = (self.bias.unsqueeze(0).expand(batch_size, *self.bias.size())) wi = torch.matmul(inp, self.weight_ih) wh = torch.matmul(h_0, self.weight_hh) i, o, g = torch.split(wh + wi + bias_batch, split_size_or_sections=self.hidden_size, dim=1) i = torch.sigmoid(i).unsqueeze(1) o = torch.sigmoid(o).unsqueeze(1) g = torch.tanh(g).unsqueeze(1) ##basic lstm start f = 1 - i c_1_basic = f*c_0.unsqueeze(1) + i*g c_1_basic = c_1_basic.squeeze(1) alpha_wi = torch.matmul(inp, self.alpha_weight_ih) alpha_wi.unsqueeze_(1) alpha_wh = torch.matmul(skip_c, self.alpha_weight_hh) alpha_bias_batch = self.alpha_bias.unsqueeze(0) alpha = torch.sigmoid(alpha_wi + alpha_wh + alpha_bias_batch) skip_mask = seq_len_to_mask(skip_count,max_len=skip_c.size()[1]) skip_mask = 1 - skip_mask skip_mask = skip_mask.unsqueeze(-1).expand(*skip_mask.size(), self.hidden_size) skip_mask = (skip_mask).float()*1e20 alpha = alpha - skip_mask alpha = torch.exp(torch.cat([i, alpha], dim=1)) alpha_sum = torch.sum(alpha, dim=1, keepdim=True) alpha = torch.div(alpha, alpha_sum) merge_i_c = torch.cat([g, skip_c], dim=1) c_1 = merge_i_c * alpha c_1 = c_1.sum(1, keepdim=True) # h_1 = o * c_1 c_1 = c_1.squeeze(1) count_select = (skip_count != 0).float().unsqueeze(-1) c_1 = c_1*count_select + c_1_basic*(1-count_select) o = o.squeeze(1) h_1 = o * torch.tanh(c_1) return h_1, c_1 class LatticeLSTMLayer_sup_back_V0(nn.Module): def __init__(self, char_input_size, word_input_size, hidden_size, left2right, bias=True,device=None,debug=False,skip_before_head=False): super().__init__() self.skip_before_head = skip_before_head self.hidden_size = hidden_size self.char_cell = MultiInputLSTMCell_V0(char_input_size, hidden_size, bias,debug) self.word_cell = WordLSTMCell_yangjie(word_input_size,hidden_size,bias,debug=self.debug) self.word_input_size = word_input_size self.left2right = left2right self.bias = bias self.device = device self.debug = debug def forward(self, inp, seq_len, skip_sources, skip_words, skip_count, init_state=None): ''' :param inp: batch * seq_len * embedding, chars :param seq_len: batch, length of chars :param skip_sources: batch * seq_len * X, 跳边的起点 :param skip_words: batch * seq_len * X * embedding, 跳边的词 :param lexicon_count: batch * seq_len, count of lexicon per example per position :param init_state: the hx of rnn :return: ''' if self.left2right: max_seq_len = max(seq_len) batch_size = inp.size(0) c_ = torch.zeros(size=[batch_size, 1, self.hidden_size], requires_grad=True).to(self.device) h_ = torch.zeros(size=[batch_size, 1, self.hidden_size], requires_grad=True).to(self.device) for i in range(max_seq_len): max_lexicon_count = max(torch.max(skip_count[:, i]).item(), 1) h_0, c_0 = h_[:, i, :], c_[:, i, :] skip_word_flat = skip_words[:, i, :max_lexicon_count].contiguous() skip_word_flat = skip_word_flat.view(batch_size*max_lexicon_count,self.word_input_size) skip_source_flat = skip_sources[:, i, :max_lexicon_count].contiguous().view(batch_size, max_lexicon_count) index_0 = torch.tensor(range(batch_size)).unsqueeze(1).expand(batch_size,max_lexicon_count) index_1 = skip_source_flat if not self.skip_before_head: c_x = c_[[index_0, index_1+1]] h_x = h_[[index_0, index_1+1]] else: c_x = c_[[index_0,index_1]] h_x = h_[[index_0,index_1]] c_x_flat = c_x.view(batch_size*max_lexicon_count,self.hidden_size) h_x_flat = h_x.view(batch_size*max_lexicon_count,self.hidden_size) c_1_flat = self.word_cell(skip_word_flat,(h_x_flat,c_x_flat)) c_1_skip = c_1_flat.view(batch_size,max_lexicon_count,self.hidden_size) h_1,c_1 = self.char_cell(inp[:,i,:],c_1_skip,skip_count[:,i],(h_0,c_0)) h_ = torch.cat([h_,h_1.unsqueeze(1)],dim=1) c_ = torch.cat([c_, c_1.unsqueeze(1)], dim=1) return h_[:,1:],c_[:,1:] else: mask_for_seq_len = seq_len_to_mask(seq_len) max_seq_len = max(seq_len) batch_size = inp.size(0) c_ = torch.zeros(size=[batch_size, 1, self.hidden_size], requires_grad=True).to(self.device) h_ = torch.zeros(size=[batch_size, 1, self.hidden_size], requires_grad=True).to(self.device) for i in reversed(range(max_seq_len)): max_lexicon_count = max(torch.max(skip_count[:, i]).item(), 1) h_0, c_0 = h_[:, 0, :], c_[:, 0, :] skip_word_flat = skip_words[:, i, :max_lexicon_count].contiguous() skip_word_flat = skip_word_flat.view(batch_size*max_lexicon_count,self.word_input_size) skip_source_flat = skip_sources[:, i, :max_lexicon_count].contiguous().view(batch_size, max_lexicon_count) index_0 = torch.tensor(range(batch_size)).unsqueeze(1).expand(batch_size,max_lexicon_count) index_1 = skip_source_flat-i if not self.skip_before_head: c_x = c_[[index_0, index_1-1]] h_x = h_[[index_0, index_1-1]] else: c_x = c_[[index_0,index_1]] h_x = h_[[index_0,index_1]] c_x_flat = c_x.view(batch_size*max_lexicon_count,self.hidden_size) h_x_flat = h_x.view(batch_size*max_lexicon_count,self.hidden_size) c_1_flat = self.word_cell(skip_word_flat,(h_x_flat,c_x_flat)) c_1_skip = c_1_flat.view(batch_size,max_lexicon_count,self.hidden_size) h_1,c_1 = self.char_cell(inp[:,i,:],c_1_skip,skip_count[:,i],(h_0,c_0)) h_1_mask = h_1.masked_fill(1-mask_for_seq_len[:,i].unsqueeze(-1),0) c_1_mask = c_1.masked_fill(1 - mask_for_seq_len[:, i].unsqueeze(-1), 0) h_ = torch.cat([h_1_mask.unsqueeze(1),h_],dim=1) c_ = torch.cat([c_1_mask.unsqueeze(1),c_], dim=1) return h_[:,:-1],c_[:,:-1] class LatticeLSTMLayer_sup_back_V1(nn.Module): # V1与V0的不同在于，V1在当前位置完全无lexicon匹配时，会采用普通的lstm计算公式， # 普通的lstm计算公式与杨杰实现的lattice lstm在lexicon数量为0时不同 def __init__(self, char_input_size, word_input_size, hidden_size, left2right, bias=True,device=None,debug=False,skip_before_head=False): super().__init__() self.debug = debug self.skip_before_head = skip_before_head self.hidden_size = hidden_size self.char_cell = MultiInputLSTMCell_V1(char_input_size, hidden_size, bias,debug) self.word_cell = WordLSTMCell_yangjie(word_input_size,hidden_size,bias,debug=self.debug) self.word_input_size = word_input_size self.left2right = left2right self.bias = bias self.device = device def forward(self, inp, seq_len, skip_sources, skip_words, skip_count, init_state=None): ''' :param inp: batch * seq_len * embedding, chars :param seq_len: batch, length of chars :param skip_sources: batch * seq_len * X, 跳边的起点 :param skip_words: batch * seq_len * X * embedding_size, 跳边的词 :param lexicon_count: batch * seq_len, lexicon_count[i,j]为第i个例子以第j个位子为结尾匹配到的词的数量 :param init_state: the hx of rnn :return: ''' if self.left2right: max_seq_len = max(seq_len) batch_size = inp.size(0) c_ = torch.zeros(size=[batch_size, 1, self.hidden_size], requires_grad=True).to(self.device) h_ = torch.zeros(size=[batch_size, 1, self.hidden_size], requires_grad=True).to(self.device) for i in range(max_seq_len): max_lexicon_count = max(torch.max(skip_count[:, i]).item(), 1) h_0, c_0 = h_[:, i, :], c_[:, i, :] #为了使rnn能够计算B*lexicon_count*embedding_size的张量，需要将其reshape成二维张量 #为了匹配pytorch的[]取址方式，需要将reshape成二维张量 skip_word_flat = skip_words[:, i, :max_lexicon_count].contiguous() skip_word_flat = skip_word_flat.view(batch_size*max_lexicon_count,self.word_input_size) skip_source_flat = skip_sources[:, i, :max_lexicon_count].contiguous().view(batch_size, max_lexicon_count) index_0 = torch.tensor(range(batch_size)).unsqueeze(1).expand(batch_size,max_lexicon_count) index_1 = skip_source_flat if not self.skip_before_head: c_x = c_[[index_0, index_1+1]] h_x = h_[[index_0, index_1+1]] else: c_x = c_[[index_0,index_1]] h_x = h_[[index_0,index_1]] c_x_flat = c_x.view(batch_size*max_lexicon_count,self.hidden_size) h_x_flat = h_x.view(batch_size*max_lexicon_count,self.hidden_size) c_1_flat = self.word_cell(skip_word_flat,(h_x_flat,c_x_flat)) c_1_skip = c_1_flat.view(batch_size,max_lexicon_count,self.hidden_size) h_1,c_1 = self.char_cell(inp[:,i,:],c_1_skip,skip_count[:,i],(h_0,c_0)) h_ = torch.cat([h_,h_1.unsqueeze(1)],dim=1) c_ = torch.cat([c_, c_1.unsqueeze(1)], dim=1) return h_[:,1:],c_[:,1:] else: mask_for_seq_len = seq_len_to_mask(seq_len) max_seq_len = max(seq_len) batch_size = inp.size(0) c_ = torch.zeros(size=[batch_size, 1, self.hidden_size], requires_grad=True).to(self.device) h_ = torch.zeros(size=[batch_size, 1, self.hidden_size], requires_grad=True).to(self.device) for i in reversed(range(max_seq_len)): max_lexicon_count = max(torch.max(skip_count[:, i]).item(), 1) h_0, c_0 = h_[:, 0, :], c_[:, 0, :] skip_word_flat = skip_words[:, i, :max_lexicon_count].contiguous() skip_word_flat = skip_word_flat.view(batch_size*max_lexicon_count,self.word_input_size) skip_source_flat = skip_sources[:, i, :max_lexicon_count].contiguous().view(batch_size, max_lexicon_count) index_0 = torch.tensor(range(batch_size)).unsqueeze(1).expand(batch_size,max_lexicon_count) index_1 = skip_source_flat-i if not self.skip_before_head: c_x = c_[[index_0, index_1-1]] h_x = h_[[index_0, index_1-1]] else: c_x = c_[[index_0,index_1]] h_x = h_[[index_0,index_1]] c_x_flat = c_x.view(batch_size*max_lexicon_count,self.hidden_size) h_x_flat = h_x.view(batch_size*max_lexicon_count,self.hidden_size) c_1_flat = self.word_cell(skip_word_flat,(h_x_flat,c_x_flat)) c_1_skip = c_1_flat.view(batch_size,max_lexicon_count,self.hidden_size) h_1,c_1 = self.char_cell(inp[:,i,:],c_1_skip,skip_count[:,i],(h_0,c_0)) h_1_mask = h_1.masked_fill(1-mask_for_seq_len[:,i].unsqueeze(-1),0) c_1_mask = c_1.masked_fill(1 - mask_for_seq_len[:, i].unsqueeze(-1), 0) h_ = torch.cat([h_1_mask.unsqueeze(1),h_],dim=1) c_ = torch.cat([c_1_mask.unsqueeze(1),c_], dim=1) return h_[:,:-1],c_[:,:-1]