# coding:utf-8 # chenjun # date:2020-04-18 import torch.nn as nn import torch import torch.nn.functional as F import numpy as np # def get_non_pad_mask(seq, PAD): # assert seq.dim() == 2 # return seq.ne(PAD).type(torch.float).unsqueeze(-1) def get_pad_mask(seq, pad_idx): return (seq == pad_idx).unsqueeze(-2) def get_subsequent_mask(seq): ''' For masking out the subsequent info. ''' sz_b, len_s = seq.size() subsequent_mask = torch.triu( torch.ones((len_s, len_s), device=seq.device, dtype=torch.uint8), diagonal=1) # 返回上三角矩阵 subsequent_mask = subsequent_mask.unsqueeze(0).expand(sz_b, -1, -1) # b x ls x ls return subsequent_mask def get_attn_key_pad_mask(seq_k, seq_q, PAD): ''' For masking out the padding part of key sequence. seq_k:src_seq seq_q:tgt_seq ''' # Expand to fit the shape of key query attention matrix. len_q = seq_q.size(1) # 目标序列 padding_mask = seq_k.eq(PAD) # 源序列 padding_mask = padding_mask.unsqueeze(1).expand(-1, len_q, -1) # b x lq x lk return padding_mask class PositionalEncoding(nn.Module): def __init__(self, d_hid, n_position=200): super(PositionalEncoding, self).__init__() # Not a parameter self.register_buffer('pos_table', self._get_sinusoid_encoding_table(n_position, d_hid)) def _get_sinusoid_encoding_table(self, n_position, d_hid): ''' Sinusoid position encoding table ''' # TODO: make it with torch instead of numpy def get_position_angle_vec(position): return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)] sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)]) sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2]) # dim 2i sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2]) # dim 2i+1 return torch.FloatTensor(sinusoid_table).unsqueeze(0) def forward(self, x): return x + self.pos_table[:, :x.size(1)].clone().detach() class ScaledDotProductAttention(nn.Module): ''' Scaled Dot-Product Attention ''' def __init__(self, temperature, attn_dropout=0.1): super(ScaledDotProductAttention, self).__init__() self.temperature = temperature self.dropout = nn.Dropout(attn_dropout) self.softmax = nn.Softmax(dim=2) def forward(self, q, k, v, mask=None): attn = torch.bmm(q, k.transpose(1, 2)) attn = attn / self.temperature if mask is not None: # print(mask.shape, attn.shape, v.shape) attn = attn.masked_fill(mask, -1e9) attn = self.softmax(attn) # 第3个维度为权重 attn = self.dropout(attn) output = torch.bmm(attn, v) return output, attn class MultiHeadAttention(nn.Module): ''' Multi-Head Attention module ''' def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1): super(MultiHeadAttention, self).__init__() self.n_head = n_head self.d_k = d_k self.d_v = d_v self.w_qs = nn.Linear(d_model, n_head * d_k) self.w_ks = nn.Linear(d_model, n_head * d_k) self.w_vs = nn.Linear(d_model, n_head * d_v) nn.init.normal_(self.w_qs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k))) nn.init.normal_(self.w_ks.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k))) nn.init.normal_(self.w_vs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_v))) self.attention = ScaledDotProductAttention(temperature=np.power(d_k, 0.5)) self.layer_norm = nn.LayerNorm(d_model) self.fc = nn.Linear(n_head * d_v, d_model) nn.init.xavier_normal_(self.fc.weight) self.dropout = nn.Dropout(dropout) def forward(self, q, k, v, mask=None): d_k, d_v, n_head = self.d_k, self.d_v, self.n_head sz_b, len_q, _ = q.size() sz_b, len_k, _ = k.size() sz_b, len_v, _ = v.size() residual = q q = self.w_qs(q).view(sz_b, len_q, n_head, d_k) # 4*21*512 ---- 4*21*8*64 k = self.w_ks(k).view(sz_b, len_k, n_head, d_k) v = self.w_vs(v).view(sz_b, len_v, n_head, d_v) q = q.permute(2, 0, 1, 3).contiguous().view(-1, len_q, d_k) # (n*b) x lq x dk k = k.permute(2, 0, 1, 3).contiguous().view(-1, len_k, d_k) # (n*b) x lk x dk v = v.permute(2, 0, 1, 3).contiguous().view(-1, len_v, d_v) # (n*b) x lv x dv mask = mask.repeat(n_head, 1, 1) if mask is not None else None # (n*b) x .. x .. output, attn = self.attention(q, k, v, mask=mask) output = output.view(n_head, sz_b, len_q, d_v) output = output.permute(1, 2, 0, 3).contiguous().view(sz_b, len_q, -1) # b x lq x (n*dv) output = self.dropout(self.fc(output)) output = self.layer_norm(output + residual) return output, attn class PositionwiseFeedForward(nn.Module): ''' A two-feed-forward-layer module ''' def __init__(self, d_in, d_hid, dropout=0.1): super(PositionwiseFeedForward, self).__init__() self.w_1 = nn.Conv1d(d_in, d_hid, 1) # position-wise self.w_2 = nn.Conv1d(d_hid, d_in, 1) # position-wise self.layer_norm = nn.LayerNorm(d_in) self.dropout = nn.Dropout(dropout) def forward(self, x): residual = x output = x.transpose(1, 2) output = self.w_2(F.relu(self.w_1(output))) output = output.transpose(1, 2) output = self.dropout(output) output = self.layer_norm(output + residual) return output class EncoderLayer(nn.Module): ''' Compose with two layers ''' def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1): super(EncoderLayer, self).__init__() self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout) self.pos_ffn = PositionwiseFeedForward(d_model, d_inner, dropout=dropout) def forward(self, enc_input, slf_attn_mask=None): enc_output, enc_slf_attn = self.slf_attn( enc_input, enc_input, enc_input, mask=slf_attn_mask) enc_output = self.pos_ffn(enc_output) return enc_output, enc_slf_attn class Torch_transformer_encoder(nn.Module): ''' use pytorch transformer for sequence learning ''' def __init__(self, d_word_vec=512, n_layers=2, n_head=8, d_model=512, dim_feedforward=1024, n_position=256): super(Torch_transformer_encoder, self).__init__() self.position_enc = PositionalEncoding(d_word_vec, n_position=n_position) encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=n_head, dim_feedforward=dim_feedforward) self.layer_norm = nn.LayerNorm(d_model) self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_layers, norm=self.layer_norm) self.dropout = nn.Dropout(p=0.1) def forward(self, cnn_feature, src_mask=None, return_attns=False): enc_slf_attn_list = [] # -- Forward enc_output = self.dropout(self.position_enc(cnn_feature)) # position embeding enc_output = self.encoder(enc_output) enc_output = self.layer_norm(enc_output) if return_attns: return enc_output, enc_slf_attn_list return enc_output, class Transforme_Encoder(nn.Module): ''' to capture the global spatial dependencies''' ''' d_word_vec: 位置编码，特征空间维度 n_layers: transformer的层数 n_head：多头数量 d_k: 64 d_v: 64 d_model: 512, d_inner: 1024 n_position: 位置编码的最大值 ''' def __init__( self, d_word_vec=512, n_layers=2, n_head=8, d_k=64, d_v=64, d_model=512, d_inner=1024, dropout=0.1, n_position=256): super().__init__() self.position_enc = PositionalEncoding(d_word_vec, n_position=n_position) self.dropout = nn.Dropout(p=dropout) self.layer_stack = nn.ModuleList([ EncoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout) for _ in range(n_layers)]) self.layer_norm = nn.LayerNorm(d_model, eps=1e-6) def forward(self, cnn_feature, src_mask, return_attns=False): enc_slf_attn_list = [] # -- Forward enc_output = self.dropout(self.position_enc(cnn_feature)) # position embeding for enc_layer in self.layer_stack: enc_output, enc_slf_attn = enc_layer(enc_output, slf_attn_mask=src_mask) enc_slf_attn_list += [enc_slf_attn] if return_attns else [] enc_output = self.layer_norm(enc_output) if return_attns: return enc_output, enc_slf_attn_list return enc_output, class PVAM(nn.Module): ''' Parallel Visual attention module 平行解码''' ''' n_dim：512，阅读顺序序列编码的空间维度 N_max_character: 25，单张图片最多有多少个字符 n_position: cnn出来之后特征的序列长度 ''' def __init__(self, n_dim=512, N_max_character=25, n_position=256): super(PVAM, self).__init__() self.character_len = N_max_character self.f0_embedding = nn.Embedding(N_max_character, n_dim) self.w0 = nn.Linear(N_max_character, n_position) self.wv = nn.Linear(n_dim, n_dim) # first linear(512,25) self.we = nn.Linear(n_dim, N_max_character) self.active = nn.Tanh() self.softmax = nn.Softmax(dim=2) def forward(self, enc_output): reading_order = torch.arange(self.character_len, dtype=torch.long, device=enc_output.device) reading_order = reading_order.unsqueeze(0).expand(enc_output.size(0), -1) # (S,) -> (B, S) reading_order_embed = self.f0_embedding(reading_order) # b,25,512 t = self.w0(reading_order_embed.permute(0,2,1)) # b,512,256 t = self.active(t.permute(0,2,1) + self.wv(enc_output)) # b,256,512 # first linear(512,25) attn = self.we(t) # b,256,25 attn = self.softmax(attn.permute(0,2,1)) # b,25,256 g_output = torch.bmm(attn, enc_output) # b,25,512 return g_output class GSRM(nn.Module): # global semantic reasoning module ''' n_dim：embed编码的特征空间维度 n_class：embedding需要用到 PAD：计算mask用到 ''' def __init__(self, n_dim=512, n_class=37, PAD=37-1, n_layers=4, n_position=25): super(GSRM, self).__init__() self.PAD = PAD self.argmax_embed = nn.Embedding(n_class, n_dim) self.transformer_units = Transforme_Encoder(n_layers=n_layers, n_position=n_position) # for global context information # self.transformer_units = Torch_transformer_encoder(n_layers=n_layers, n_position=n_position) def forward(self, e_out): ''' e_out: b,25,37 | the output from PVAM3 ''' e_argmax = e_out.argmax(dim=-1) # b, 25 e = self.argmax_embed(e_argmax) # b,25,512 e_mask = get_pad_mask(e_argmax, self.PAD) # b,25,1 s = self.transformer_units(e, None) # b,25,512 return s class SRN_Decoder(nn.Module): # the wrapper of decoder layers ''' n_dim: 特征空间维度 n_class：字符种类 N_max_character: 单张图最多只25个字符 n_position：cnn输出的特征序列长度 整个有三个部分的输出 ''' def __init__(self, n_dim=512, n_class=37, N_max_character=25, n_position=256, GSRM_layer=4 ): super(SRN_Decoder, self).__init__() self.pvam = PVAM(N_max_character=N_max_character, n_position=n_position) self.w_e = nn.Linear(n_dim, n_class) # output layer self.GSRM = GSRM(n_class=n_class, PAD=n_class-1, n_dim=n_dim, n_position=N_max_character, n_layers=GSRM_layer) self.w_s = nn.Linear(n_dim, n_class) # output layer self.w_f = nn.Linear(n_dim, n_class) # output layer def forward(self, cnn_feature ): '''cnn_feature: b,256,512 | the output from cnn''' g_output = self.pvam(cnn_feature) # b,25,512 e_out = self.w_e(g_output) # b,25,37 ----> cross entropy loss | 第一个输出 s = self.GSRM(e_out)[0] # b,25,512 s_out = self.w_s(s) # b,25,37f # TODO:change the add to gated unit f = g_output + s # b,25,512 f_out = self.w_f(f) return e_out, s_out, f_out def cal_performance(preds, gold, mask=None, smoothing='1'): ''' Apply label smoothing if needed ''' loss = 0. n_correct = 0 weights = [1.0, 0.15, 2.0] for ori_pred, weight in zip(preds, weights): pred = ori_pred.view(-1, ori_pred.shape[-1]) # debug show t_gold = gold.view(ori_pred.shape[0], -1) t_pred_index = ori_pred.max(2)[1] mask = mask.view(-1) non_pad_mask = mask.ne(0) if mask is not None else None tloss = cal_loss(pred, gold, non_pad_mask, smoothing) if torch.isnan(tloss): print('have nan loss') continue else: loss += tloss * weight pred = pred.max(1)[1] gold = gold.contiguous().view(-1) n_correct = pred.eq(gold) n_correct = n_correct.masked_select(non_pad_mask).sum().item() if mask is not None else None return loss, n_correct def cal_loss(pred, gold, mask, smoothing): ''' Calculate cross entropy loss, apply label smoothing if needed. ''' gold = gold.contiguous().view(-1) if smoothing=='0': eps = 0.1 n_class = pred.size(1) one_hot = torch.zeros_like(pred).scatter(1, gold.view(-1, 1), 1) one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (n_class - 1) log_prb = F.log_softmax(pred, dim=1) non_pad_mask = gold.ne(0) loss = -(one_hot * log_prb).sum(dim=1) loss = loss.masked_select(non_pad_mask).sum() # average later elif smoothing == '1': if mask is not None: loss = F.cross_entropy(pred, gold, reduction='none') loss = loss.masked_select(mask) loss = loss.sum() / mask.sum() else: loss = F.cross_entropy(pred, gold) else: # loss = F.cross_entropy(pred, gold, ignore_index=PAD) loss = F.cross_entropy(pred, gold) return loss def cal_performance2(preds, gold, PAD, smoothing='1'): ''' Apply label smoothing if needed ''' loss = 0. n_correct = 0 weights = [1.0, 0.15, 2.0] for ori_pred, weight in zip(preds, weights): pred = ori_pred.view(-1, ori_pred.shape[-1]) # debug show t_gold = gold.view(ori_pred.shape[0], -1) t_pred_index = ori_pred.max(2)[1] tloss = cal_loss2(pred, gold, PAD, smoothing=smoothing) if torch.isnan(tloss): print('have nan loss') continue else: loss += tloss * weight pred = pred.max(1)[1] gold = gold.contiguous().view(-1) n_correct = pred.eq(gold) non_pad_mask = gold.ne(PAD) n_correct = n_correct.masked_select(non_pad_mask).sum().item() return loss, n_correct def cal_loss2(pred, gold, PAD, smoothing='1'): ''' Calculate cross entropy loss, apply label smoothing if needed. ''' gold = gold.contiguous().view(-1) if smoothing=='0': eps = 0.1 n_class = pred.size(1) one_hot = torch.zeros_like(pred).scatter(1, gold.view(-1, 1), 1) one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (n_class - 1) log_prb = F.log_softmax(pred, dim=1) non_pad_mask = gold.ne(0) loss = -(one_hot * log_prb).sum(dim=1) loss = loss.masked_select(non_pad_mask).sum() # average later elif smoothing == '1': loss = F.cross_entropy(pred, gold, ignore_index=PAD) else: # loss = F.cross_entropy(pred, gold, ignore_index=PAD) loss = F.cross_entropy(pred, gold) return loss if __name__=='__main__': cnn_feature = torch.rand((2,256,512)) model1 = Transforme_Encoder() image = model1(cnn_feature,src_mask=None)[0] model = SRN_Decoder(N_max_character=30) outs = model(image) for out in outs: print(out.shape) # image = torch.rand((4,3,32,60)) # tgt_seq = torch.tensor([[ 2, 24, 2176, 882, 2480, 612, 1525, 480, 875, 147, 1700, 715, # 1465, 3], # [ 2, 369, 1781, 882, 703, 879, 2855, 2415, 502, 1154, 833, 1465, # 3, 0], # [ 2, 2943, 334, 328, 480, 330, 1644, 1449, 163, 147, 1823, 1184, # 1465, 3], # [ 2, 24, 396, 480, 703, 1646, 897, 1711, 1508, 703, 2321, 147, # 642, 1465]], device='cuda:0') # tgt_pos = torch.tensor([[ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14], # [ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 0], # [ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14], # [ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]], # device='cuda:0') # src_seq = torch.tensor([[ 2, 598, 2088, 822, 2802, 1156, 157, 1099, 1000, 598, 1707, 1345, # 3, 0, 0, 0], # [ 2, 598, 2348, 822, 598, 1222, 471, 948, 986, 423, 1345, 3, # 0, 0, 0, 0], # [ 2, 2437, 2470, 901, 2473, 598, 1735, 84, 1, 2277, 1979, 499, # 962, 1345, 3, 0], # [ 2, 598, 186, 1904, 598, 868, 1339, 1604, 84, 598, 608, 1728, # 1345, 3, 0, 0]], device='cuda:0') # device = torch.device('cuda') # image = image.cuda() # transformer = Transformer() # transformer = transformer.to(device) # transformer.train() # out = transformer(image, tgt_seq, tgt_pos, src_seq) # gold = tgt_seq[:, 1:] # 从第二列开始 # # backward # loss, n_correct = cal_performance(out, gold, smoothing=True) # print(loss, n_correct) # a = 1