modules_srn/SRN_modules.py

# 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