First model version

app.py

@@ -10,7 +10,7 @@ from demo import get_model, preprocess, postprocess, load
 from utils import Config, Logger, CharsetMapper
 
 def process_image(image):
-  config = Config('configs/train_abinet.yaml')
+  config = Config('configs/rec/train_abinet.yaml')
   config.model_vision_checkpoint = None
   model = get_model(config)
   model = load(model, 'workdir/train-abinet/best-train-abinet.pth')

modules/attention.py

@@ -1,97 +0,0 @@
-import torch
-import torch.nn as nn
-from .transformer import PositionalEncoding
-
-class Attention(nn.Module):
-    def __init__(self, in_channels=512, max_length=25, n_feature=256):
-        super().__init__()
-        self.max_length = max_length
-
-        self.f0_embedding = nn.Embedding(max_length, in_channels)
-        self.w0 = nn.Linear(max_length, n_feature)
-        self.wv = nn.Linear(in_channels, in_channels)
-        self.we = nn.Linear(in_channels, max_length)
-
-        self.active = nn.Tanh()
-        self.softmax = nn.Softmax(dim=2)
-
-    def forward(self, enc_output):
-        enc_output = enc_output.permute(0, 2, 3, 1).flatten(1, 2)
-        reading_order = torch.arange(self.max_length, 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
-
-        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, attn.view(*attn.shape[:2], 8, 32)
-
-
-def encoder_layer(in_c, out_c, k=3, s=2, p=1):
-    return nn.Sequential(nn.Conv2d(in_c, out_c, k, s, p),
-                         nn.BatchNorm2d(out_c),
-                         nn.ReLU(True))
-
-def decoder_layer(in_c, out_c, k=3, s=1, p=1, mode='nearest', scale_factor=None, size=None):
-    align_corners = None if mode=='nearest' else True
-    return nn.Sequential(nn.Upsample(size=size, scale_factor=scale_factor, 
-                                     mode=mode, align_corners=align_corners),
-                         nn.Conv2d(in_c, out_c, k, s, p),
-                         nn.BatchNorm2d(out_c),
-                         nn.ReLU(True))
-
-
-class PositionAttention(nn.Module):
-    def __init__(self, max_length, in_channels=512, num_channels=64, 
-                 h=8, w=32, mode='nearest', **kwargs):
-        super().__init__()
-        self.max_length = max_length
-        self.k_encoder = nn.Sequential(
-            encoder_layer(in_channels, num_channels, s=(1, 2)),
-            encoder_layer(num_channels, num_channels, s=(2, 2)),
-            encoder_layer(num_channels, num_channels, s=(2, 2)),
-            encoder_layer(num_channels, num_channels, s=(2, 2))
-        )
-        self.k_decoder = nn.Sequential(
-            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
-            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
-            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
-            decoder_layer(num_channels, in_channels, size=(h, w), mode=mode)
-        )
-
-        self.pos_encoder = PositionalEncoding(in_channels, dropout=0, max_len=max_length)
-        self.project = nn.Linear(in_channels, in_channels)
-
-    def forward(self, x):
-        N, E, H, W = x.size()
-        k, v = x, x  # (N, E, H, W)
-
-        # calculate key vector
-        features = []
-        for i in range(0, len(self.k_encoder)):
-            k = self.k_encoder[i](k)
-            features.append(k)
-        for i in range(0, len(self.k_decoder) - 1):
-            k = self.k_decoder[i](k)
-            k = k + features[len(self.k_decoder) - 2 - i]
-        k = self.k_decoder[-1](k)
-
-        # calculate query vector
-        # TODO q=f(q,k)
-        zeros = x.new_zeros((self.max_length, N, E))  # (T, N, E)
-        q = self.pos_encoder(zeros)  # (T, N, E)
-        q = q.permute(1, 0, 2)  # (N, T, E)
-        q = self.project(q)  # (N, T, E)
-        
-        # calculate attention
-        attn_scores = torch.bmm(q, k.flatten(2, 3))  # (N, T, (H*W))
-        attn_scores = attn_scores / (E ** 0.5)
-        attn_scores = torch.softmax(attn_scores, dim=-1)
-
-        v = v.permute(0, 2, 3, 1).view(N, -1, E)  # (N, (H*W), E)
-        attn_vecs = torch.bmm(attn_scores, v)  # (N, T, E)
-
-        return attn_vecs, attn_scores.view(N, -1, H, W)

modules/backbone.py

@@ -1,36 +0,0 @@
-import torch
-import torch.nn as nn
-from fastai.vision import *
-
-from modules.model import _default_tfmer_cfg
-from modules.resnet import resnet45
-from modules.transformer import (PositionalEncoding,
-                                 TransformerEncoder,
-                                 TransformerEncoderLayer)
-
-
-class ResTranformer(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.resnet = resnet45()
-
-        self.d_model = ifnone(config.model_vision_d_model, _default_tfmer_cfg['d_model'])
-        nhead = ifnone(config.model_vision_nhead, _default_tfmer_cfg['nhead'])
-        d_inner = ifnone(config.model_vision_d_inner, _default_tfmer_cfg['d_inner'])
-        dropout = ifnone(config.model_vision_dropout, _default_tfmer_cfg['dropout'])
-        activation = ifnone(config.model_vision_activation, _default_tfmer_cfg['activation'])
-        num_layers = ifnone(config.model_vision_backbone_ln, 2)
-
-        self.pos_encoder = PositionalEncoding(self.d_model, max_len=8*32)
-        encoder_layer = TransformerEncoderLayer(d_model=self.d_model, nhead=nhead, 
-                dim_feedforward=d_inner, dropout=dropout, activation=activation)
-        self.transformer = TransformerEncoder(encoder_layer, num_layers)
-
-    def forward(self, images):
-        feature = self.resnet(images)
-        n, c, h, w = feature.shape
-        feature = feature.view(n, c, -1).permute(2, 0, 1)
-        feature = self.pos_encoder(feature)
-        feature = self.transformer(feature)
-        feature = feature.permute(1, 2, 0).view(n, c, h, w)
-        return feature

modules/model.py

@@ -1,50 +0,0 @@
-import torch
-import torch.nn as nn
-
-from utils import CharsetMapper
-
-
-_default_tfmer_cfg = dict(d_model=512, nhead=8, d_inner=2048, # 1024
-                          dropout=0.1, activation='relu')
-
-class Model(nn.Module):
-
-    def __init__(self, config):
-        super().__init__()
-        self.max_length = config.dataset_max_length + 1
-        self.charset = CharsetMapper(config.dataset_charset_path, max_length=self.max_length)
-    
-    def load(self, source, device=None, strict=True):
-        state = torch.load(source, map_location=device)
-        self.load_state_dict(state['model'], strict=strict)
-
-    def _get_length(self, logit, dim=-1):
-        """ Greed decoder to obtain length from logit"""
-        out = (logit.argmax(dim=-1) == self.charset.null_label)
-        abn = out.any(dim)
-        out = ((out.cumsum(dim) == 1) & out).max(dim)[1]
-        out = out + 1  # additional end token
-        out = torch.where(abn, out, out.new_tensor(logit.shape[1]))
-        return out
-
-    @staticmethod
-    def _get_padding_mask(length, max_length):
-        length = length.unsqueeze(-1)
-        grid = torch.arange(0, max_length, device=length.device).unsqueeze(0)
-        return grid >= length
-
-    @staticmethod
-    def _get_square_subsequent_mask(sz, device, diagonal=0, fw=True):
-        r"""Generate a square mask for the sequence. The masked positions are filled with float('-inf').
-            Unmasked positions are filled with float(0.0).
-        """
-        mask = (torch.triu(torch.ones(sz, sz, device=device), diagonal=diagonal) == 1)
-        if fw: mask = mask.transpose(0, 1)
-        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
-        return mask
-
-    @staticmethod
-    def _get_location_mask(sz, device=None):
-        mask = torch.eye(sz, device=device)
-        mask = mask.float().masked_fill(mask == 1, float('-inf'))
-        return mask

modules/model_abinet.py

@@ -1,30 +0,0 @@
-import torch
-import torch.nn as nn
-from fastai.vision import *
-
-from .model_vision import BaseVision
-from .model_language import BCNLanguage
-from .model_alignment import BaseAlignment
-
-
-class ABINetModel(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.use_alignment = ifnone(config.model_use_alignment, True)
-        self.max_length = config.dataset_max_length + 1  # additional stop token
-        self.vision = BaseVision(config)
-        self.language = BCNLanguage(config)
-        if self.use_alignment: self.alignment = BaseAlignment(config)
-
-    def forward(self, images, *args):
-        v_res = self.vision(images)
-        v_tokens = torch.softmax(v_res['logits'], dim=-1)
-        v_lengths = v_res['pt_lengths'].clamp_(2, self.max_length)  # TODO:move to langauge model
-
-        l_res = self.language(v_tokens, v_lengths)
-        if not self.use_alignment:
-            return l_res, v_res
-        l_feature, v_feature = l_res['feature'], v_res['feature']
-        
-        a_res = self.alignment(l_feature, v_feature)
-        return a_res, l_res, v_res

modules/model_abinet_iter.py

@@ -1,34 +0,0 @@
-import torch
-import torch.nn as nn
-from fastai.vision import *
-
-from .model_vision import BaseVision
-from .model_language import BCNLanguage
-from .model_alignment import BaseAlignment
-
-
-class ABINetIterModel(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.iter_size = ifnone(config.model_iter_size, 1)
-        self.max_length = config.dataset_max_length + 1  # additional stop token
-        self.vision = BaseVision(config)
-        self.language = BCNLanguage(config)
-        self.alignment = BaseAlignment(config)
-
-    def forward(self, images, *args):
-        v_res = self.vision(images)
-        a_res = v_res
-        all_l_res, all_a_res = [], []
-        for _ in range(self.iter_size):
-            tokens = torch.softmax(a_res['logits'], dim=-1)
-            lengths = a_res['pt_lengths']
-            lengths.clamp_(2, self.max_length)  # TODO:move to langauge model
-            l_res = self.language(tokens, lengths)
-            all_l_res.append(l_res)
-            a_res = self.alignment(l_res['feature'], v_res['feature'])
-            all_a_res.append(a_res)
-        if self.training:
-            return all_a_res, all_l_res, v_res
-        else:
-            return a_res, all_l_res[-1], v_res

modules/model_alignment.py

@@ -1,34 +0,0 @@
-import torch
-import torch.nn as nn
-from fastai.vision import *
-
-from modules.model import Model, _default_tfmer_cfg
-
-
-class BaseAlignment(Model):
-    def __init__(self, config):
-        super().__init__(config)
-        d_model = ifnone(config.model_alignment_d_model, _default_tfmer_cfg['d_model'])
-
-        self.loss_weight = ifnone(config.model_alignment_loss_weight, 1.0)
-        self.max_length = config.dataset_max_length + 1  # additional stop token
-        self.w_att = nn.Linear(2 * d_model, d_model)
-        self.cls = nn.Linear(d_model, self.charset.num_classes)
-
-    def forward(self, l_feature, v_feature):
-        """
-        Args:
-            l_feature: (N, T, E) where T is length, N is batch size and d is dim of model
-            v_feature: (N, T, E) shape the same as l_feature 
-            l_lengths: (N,)
-            v_lengths: (N,)
-        """
-        f = torch.cat((l_feature, v_feature), dim=2)
-        f_att = torch.sigmoid(self.w_att(f))
-        output = f_att * v_feature + (1 - f_att) * l_feature
-
-        logits = self.cls(output)  # (N, T, C)
-        pt_lengths = self._get_length(logits)
-
-        return {'logits': logits, 'pt_lengths': pt_lengths, 'loss_weight':self.loss_weight,
-                'name': 'alignment'}

modules/model_language.py

@@ -1,67 +0,0 @@
-import logging
-import torch.nn as nn
-from fastai.vision import *
-
-from modules.model import _default_tfmer_cfg
-from modules.model import Model
-from modules.transformer import (PositionalEncoding, 
-                                 TransformerDecoder,
-                                 TransformerDecoderLayer)
-
-
-class BCNLanguage(Model):
-    def __init__(self, config):
-        super().__init__(config)
-        d_model = ifnone(config.model_language_d_model, _default_tfmer_cfg['d_model'])
-        nhead = ifnone(config.model_language_nhead, _default_tfmer_cfg['nhead'])
-        d_inner = ifnone(config.model_language_d_inner, _default_tfmer_cfg['d_inner'])
-        dropout = ifnone(config.model_language_dropout, _default_tfmer_cfg['dropout'])
-        activation = ifnone(config.model_language_activation, _default_tfmer_cfg['activation'])
-        num_layers = ifnone(config.model_language_num_layers, 4)
-        self.d_model = d_model
-        self.detach = ifnone(config.model_language_detach, True)
-        self.use_self_attn = ifnone(config.model_language_use_self_attn, False)
-        self.loss_weight = ifnone(config.model_language_loss_weight, 1.0)
-        self.max_length = config.dataset_max_length + 1  # additional stop token
-        self.debug = ifnone(config.global_debug, False)
-
-        self.proj = nn.Linear(self.charset.num_classes, d_model, False)
-        self.token_encoder = PositionalEncoding(d_model, max_len=self.max_length)
-        self.pos_encoder = PositionalEncoding(d_model, dropout=0, max_len=self.max_length)
-        decoder_layer = TransformerDecoderLayer(d_model, nhead, d_inner, dropout, 
-                activation, self_attn=self.use_self_attn, debug=self.debug)
-        self.model = TransformerDecoder(decoder_layer, num_layers)
-
-        self.cls = nn.Linear(d_model, self.charset.num_classes)
-
-        if config.model_language_checkpoint is not None:
-            logging.info(f'Read language model from {config.model_language_checkpoint}.')
-            self.load(config.model_language_checkpoint)
-
-    def forward(self, tokens, lengths):
-        """
-        Args:
-            tokens: (N, T, C) where T is length, N is batch size and C is classes number
-            lengths: (N,)
-        """
-        if self.detach: tokens = tokens.detach()
-        embed = self.proj(tokens)  # (N, T, E)
-        embed = embed.permute(1, 0, 2)  # (T, N, E)
-        embed = self.token_encoder(embed)  # (T, N, E)
-        padding_mask = self._get_padding_mask(lengths, self.max_length)
-
-        zeros = embed.new_zeros(*embed.shape)
-        qeury = self.pos_encoder(zeros)
-        location_mask = self._get_location_mask(self.max_length, tokens.device)
-        output = self.model(qeury, embed,
-                tgt_key_padding_mask=padding_mask,
-                memory_mask=location_mask,
-                memory_key_padding_mask=padding_mask)  # (T, N, E)
-        output = output.permute(1, 0, 2)  # (N, T, E)
-
-        logits = self.cls(output)  # (N, T, C)
-        pt_lengths = self._get_length(logits)
-
-        res =  {'feature': output, 'logits': logits, 'pt_lengths': pt_lengths,
-                'loss_weight':self.loss_weight, 'name': 'language'}
-        return res

modules/model_vision.py

@@ -1,47 +0,0 @@
-import logging
-import torch.nn as nn
-from fastai.vision import *
-
-from modules.attention import *
-from modules.backbone import ResTranformer
-from modules.model import Model
-from modules.resnet import resnet45
-
-
-class BaseVision(Model):
-    def __init__(self, config):
-        super().__init__(config)
-        self.loss_weight = ifnone(config.model_vision_loss_weight, 1.0)
-        self.out_channels = ifnone(config.model_vision_d_model, 512)
-
-        if config.model_vision_backbone == 'transformer':
-            self.backbone = ResTranformer(config)
-        else: self.backbone = resnet45()
-        
-        if config.model_vision_attention == 'position':
-            mode = ifnone(config.model_vision_attention_mode, 'nearest')
-            self.attention = PositionAttention(
-                max_length=config.dataset_max_length + 1,  # additional stop token
-                mode=mode,
-            )
-        elif config.model_vision_attention == 'attention':
-            self.attention = Attention(
-                max_length=config.dataset_max_length + 1,  # additional stop token
-                n_feature=8*32,
-            )
-        else:
-            raise Exception(f'{config.model_vision_attention} is not valid.')
-        self.cls = nn.Linear(self.out_channels, self.charset.num_classes)
-
-        if config.model_vision_checkpoint is not None:
-            logging.info(f'Read vision model from {config.model_vision_checkpoint}.')
-            self.load(config.model_vision_checkpoint)
-
-    def forward(self, images, *args):
-        features = self.backbone(images)  # (N, E, H, W)
-        attn_vecs, attn_scores = self.attention(features)  # (N, T, E), (N, T, H, W)
-        logits = self.cls(attn_vecs) # (N, T, C)
-        pt_lengths = self._get_length(logits)
-
-        return {'feature': attn_vecs, 'logits': logits, 'pt_lengths': pt_lengths,
-                'attn_scores': attn_scores, 'loss_weight':self.loss_weight, 'name': 'vision'}

modules/resnet.py

@@ -1,104 +0,0 @@
-import math
-
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.utils.model_zoo as model_zoo
-
-
-def conv1x1(in_planes, out_planes, stride=1):
-    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
-
-
-def conv3x3(in_planes, out_planes, stride=1):
-    "3x3 convolution with padding"
-    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
-                     padding=1, bias=False)
-
-
-class BasicBlock(nn.Module):
-    expansion = 1
-
-    def __init__(self, inplanes, planes, stride=1, downsample=None):
-        super(BasicBlock, self).__init__()
-        self.conv1 = conv1x1(inplanes, planes)
-        self.bn1 = nn.BatchNorm2d(planes)
-        self.relu = nn.ReLU(inplace=True)
-        self.conv2 = conv3x3(planes, planes, stride)
-        self.bn2 = nn.BatchNorm2d(planes)
-        self.downsample = downsample
-        self.stride = stride
-
-    def forward(self, x):
-        residual = x
-
-        out = self.conv1(x)
-        out = self.bn1(out)
-        out = self.relu(out)
-
-        out = self.conv2(out)
-        out = self.bn2(out)
-
-        if self.downsample is not None:
-            residual = self.downsample(x)
-
-        out += residual
-        out = self.relu(out)
-
-        return out
-
-
-class ResNet(nn.Module):
-
-    def __init__(self, block, layers):
-        self.inplanes = 32
-        super(ResNet, self).__init__()
-        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1,
-                               bias=False)
-        self.bn1 = nn.BatchNorm2d(32)
-        self.relu = nn.ReLU(inplace=True)
-
-        self.layer1 = self._make_layer(block, 32, layers[0], stride=2)
-        self.layer2 = self._make_layer(block, 64, layers[1], stride=1)
-        self.layer3 = self._make_layer(block, 128, layers[2], stride=2)
-        self.layer4 = self._make_layer(block, 256, layers[3], stride=1)
-        self.layer5 = self._make_layer(block, 512, layers[4], stride=1)
-
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
-                m.weight.data.normal_(0, math.sqrt(2. / n))
-            elif isinstance(m, nn.BatchNorm2d):
-                m.weight.data.fill_(1)
-                m.bias.data.zero_()
-
-    def _make_layer(self, block, planes, blocks, stride=1):
-        downsample = None
-        if stride != 1 or self.inplanes != planes * block.expansion:
-            downsample = nn.Sequential(
-                nn.Conv2d(self.inplanes, planes * block.expansion,
-                          kernel_size=1, stride=stride, bias=False),
-                nn.BatchNorm2d(planes * block.expansion),
-            )
-
-        layers = []
-        layers.append(block(self.inplanes, planes, stride, downsample))
-        self.inplanes = planes * block.expansion
-        for i in range(1, blocks):
-            layers.append(block(self.inplanes, planes))
-
-        return nn.Sequential(*layers)
-
-    def forward(self, x):
-        x = self.conv1(x)
-        x = self.bn1(x)
-        x = self.relu(x)
-        x = self.layer1(x)
-        x = self.layer2(x)
-        x = self.layer3(x)
-        x = self.layer4(x)
-        x = self.layer5(x)
-        return x
-
-
-def resnet45():
-    return ResNet(BasicBlock, [3, 4, 6, 6, 3])

modules/transformer.py

@@ -1,901 +0,0 @@
-# pytorch 1.5.0
-import copy
-import math
-import warnings
-from typing import Optional
-
-import torch
-import torch.nn as nn
-from torch import Tensor
-from torch.nn import Dropout, LayerNorm, Linear, Module, ModuleList, Parameter
-from torch.nn import functional as F
-from torch.nn.init import constant_, xavier_uniform_
-
-
-def multi_head_attention_forward(query,                           # type: Tensor
-                                 key,                             # type: Tensor
-                                 value,                           # type: Tensor
-                                 embed_dim_to_check,              # type: int
-                                 num_heads,                       # type: int
-                                 in_proj_weight,                  # type: Tensor
-                                 in_proj_bias,                    # type: Tensor
-                                 bias_k,                          # type: Optional[Tensor]
-                                 bias_v,                          # type: Optional[Tensor]
-                                 add_zero_attn,                   # type: bool
-                                 dropout_p,                       # type: float
-                                 out_proj_weight,                 # type: Tensor
-                                 out_proj_bias,                   # type: Tensor
-                                 training=True,                   # type: bool
-                                 key_padding_mask=None,           # type: Optional[Tensor]
-                                 need_weights=True,               # type: bool
-                                 attn_mask=None,                  # type: Optional[Tensor]
-                                 use_separate_proj_weight=False,  # type: bool
-                                 q_proj_weight=None,              # type: Optional[Tensor]
-                                 k_proj_weight=None,              # type: Optional[Tensor]
-                                 v_proj_weight=None,              # type: Optional[Tensor]
-                                 static_k=None,                   # type: Optional[Tensor]
-                                 static_v=None                    # type: Optional[Tensor]
-                                 ):
-    # type: (...) -> Tuple[Tensor, Optional[Tensor]]
-    r"""
-    Args:
-        query, key, value: map a query and a set of key-value pairs to an output.
-            See "Attention Is All You Need" for more details.
-        embed_dim_to_check: total dimension of the model.
-        num_heads: parallel attention heads.
-        in_proj_weight, in_proj_bias: input projection weight and bias.
-        bias_k, bias_v: bias of the key and value sequences to be added at dim=0.
-        add_zero_attn: add a new batch of zeros to the key and
-                       value sequences at dim=1.
-        dropout_p: probability of an element to be zeroed.
-        out_proj_weight, out_proj_bias: the output projection weight and bias.
-        training: apply dropout if is ``True``.
-        key_padding_mask: if provided, specified padding elements in the key will
-            be ignored by the attention. This is an binary mask. When the value is True,
-            the corresponding value on the attention layer will be filled with -inf.
-        need_weights: output attn_output_weights.
-        attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
-            the batches while a 3D mask allows to specify a different mask for the entries of each batch.
-        use_separate_proj_weight: the function accept the proj. weights for query, key,
-            and value in different forms. If false, in_proj_weight will be used, which is
-            a combination of q_proj_weight, k_proj_weight, v_proj_weight.
-        q_proj_weight, k_proj_weight, v_proj_weight, in_proj_bias: input projection weight and bias.
-        static_k, static_v: static key and value used for attention operators.
-    Shape:
-        Inputs:
-        - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
-          the embedding dimension.
-        - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
-          the embedding dimension.
-        - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
-          the embedding dimension.
-        - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
-          If a ByteTensor is provided, the non-zero positions will be ignored while the zero positions
-          will be unchanged. If a BoolTensor is provided, the positions with the
-          value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
-        - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
-          3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
-          S is the source sequence length. attn_mask ensures that position i is allowed to attend the unmasked
-          positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
-          while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
-          are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
-          is provided, it will be added to the attention weight.
-        - static_k: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
-          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
-        - static_v: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
-          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
-        Outputs:
-        - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
-          E is the embedding dimension.
-        - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
-          L is the target sequence length, S is the source sequence length.
-    """
-    # if not torch.jit.is_scripting():
-    #     tens_ops = (query, key, value, in_proj_weight, in_proj_bias, bias_k, bias_v,
-    #                 out_proj_weight, out_proj_bias)
-    #     if any([type(t) is not Tensor for t in tens_ops]) and has_torch_function(tens_ops):
-    #         return handle_torch_function(
-    #             multi_head_attention_forward, tens_ops, query, key, value,
-    #             embed_dim_to_check, num_heads, in_proj_weight, in_proj_bias,
-    #             bias_k, bias_v, add_zero_attn, dropout_p, out_proj_weight,
-    #             out_proj_bias, training=training, key_padding_mask=key_padding_mask,
-    #             need_weights=need_weights, attn_mask=attn_mask,
-    #             use_separate_proj_weight=use_separate_proj_weight,
-    #             q_proj_weight=q_proj_weight, k_proj_weight=k_proj_weight,
-    #             v_proj_weight=v_proj_weight, static_k=static_k, static_v=static_v)
-    tgt_len, bsz, embed_dim = query.size()
-    assert embed_dim == embed_dim_to_check
-    assert key.size() == value.size()
-
-    head_dim = embed_dim // num_heads
-    assert head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"
-    scaling = float(head_dim) ** -0.5
-
-    if not use_separate_proj_weight:
-        if torch.equal(query, key) and torch.equal(key, value):
-            # self-attention
-            q, k, v = F.linear(query, in_proj_weight, in_proj_bias).chunk(3, dim=-1)
-
-        elif torch.equal(key, value):
-            # encoder-decoder attention
-            # This is inline in_proj function with in_proj_weight and in_proj_bias
-            _b = in_proj_bias
-            _start = 0
-            _end = embed_dim
-            _w = in_proj_weight[_start:_end, :]
-            if _b is not None:
-                _b = _b[_start:_end]
-            q = F.linear(query, _w, _b)
-
-            if key is None:
-                assert value is None
-                k = None
-                v = None
-            else:
-
-                # This is inline in_proj function with in_proj_weight and in_proj_bias
-                _b = in_proj_bias
-                _start = embed_dim
-                _end = None
-                _w = in_proj_weight[_start:, :]
-                if _b is not None:
-                    _b = _b[_start:]
-                k, v = F.linear(key, _w, _b).chunk(2, dim=-1)
-
-        else:
-            # This is inline in_proj function with in_proj_weight and in_proj_bias
-            _b = in_proj_bias
-            _start = 0
-            _end = embed_dim
-            _w = in_proj_weight[_start:_end, :]
-            if _b is not None:
-                _b = _b[_start:_end]
-            q = F.linear(query, _w, _b)
-
-            # This is inline in_proj function with in_proj_weight and in_proj_bias
-            _b = in_proj_bias
-            _start = embed_dim
-            _end = embed_dim * 2
-            _w = in_proj_weight[_start:_end, :]
-            if _b is not None:
-                _b = _b[_start:_end]
-            k = F.linear(key, _w, _b)
-
-            # This is inline in_proj function with in_proj_weight and in_proj_bias
-            _b = in_proj_bias
-            _start = embed_dim * 2
-            _end = None
-            _w = in_proj_weight[_start:, :]
-            if _b is not None:
-                _b = _b[_start:]
-            v = F.linear(value, _w, _b)
-    else:
-        q_proj_weight_non_opt = torch.jit._unwrap_optional(q_proj_weight)
-        len1, len2 = q_proj_weight_non_opt.size()
-        assert len1 == embed_dim and len2 == query.size(-1)
-
-        k_proj_weight_non_opt = torch.jit._unwrap_optional(k_proj_weight)
-        len1, len2 = k_proj_weight_non_opt.size()
-        assert len1 == embed_dim and len2 == key.size(-1)
-
-        v_proj_weight_non_opt = torch.jit._unwrap_optional(v_proj_weight)
-        len1, len2 = v_proj_weight_non_opt.size()
-        assert len1 == embed_dim and len2 == value.size(-1)
-
-        if in_proj_bias is not None:
-            q = F.linear(query, q_proj_weight_non_opt, in_proj_bias[0:embed_dim])
-            k = F.linear(key, k_proj_weight_non_opt, in_proj_bias[embed_dim:(embed_dim * 2)])
-            v = F.linear(value, v_proj_weight_non_opt, in_proj_bias[(embed_dim * 2):])
-        else:
-            q = F.linear(query, q_proj_weight_non_opt, in_proj_bias)
-            k = F.linear(key, k_proj_weight_non_opt, in_proj_bias)
-            v = F.linear(value, v_proj_weight_non_opt, in_proj_bias)
-    q = q * scaling
-
-    if attn_mask is not None:
-        assert attn_mask.dtype == torch.float32 or attn_mask.dtype == torch.float64 or \
-            attn_mask.dtype == torch.float16 or attn_mask.dtype == torch.uint8 or attn_mask.dtype == torch.bool, \
-            'Only float, byte, and bool types are supported for attn_mask, not {}'.format(attn_mask.dtype)
-        if attn_mask.dtype == torch.uint8:
-            warnings.warn("Byte tensor for attn_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
-            attn_mask = attn_mask.to(torch.bool)
-
-        if attn_mask.dim() == 2:
-            attn_mask = attn_mask.unsqueeze(0)
-            if list(attn_mask.size()) != [1, query.size(0), key.size(0)]:
-                raise RuntimeError('The size of the 2D attn_mask is not correct.')
-        elif attn_mask.dim() == 3:
-            if list(attn_mask.size()) != [bsz * num_heads, query.size(0), key.size(0)]:
-                raise RuntimeError('The size of the 3D attn_mask is not correct.')
-        else:
-            raise RuntimeError("attn_mask's dimension {} is not supported".format(attn_mask.dim()))
-        # attn_mask's dim is 3 now.
-
-    # # convert ByteTensor key_padding_mask to bool
-    # if key_padding_mask is not None and key_padding_mask.dtype == torch.uint8:
-    #     warnings.warn("Byte tensor for key_padding_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
-    #     key_padding_mask = key_padding_mask.to(torch.bool)
-
-    if bias_k is not None and bias_v is not None:
-        if static_k is None and static_v is None:
-            k = torch.cat([k, bias_k.repeat(1, bsz, 1)])
-            v = torch.cat([v, bias_v.repeat(1, bsz, 1)])
-            if attn_mask is not None:
-                attn_mask = pad(attn_mask, (0, 1))
-            if key_padding_mask is not None:
-                key_padding_mask = pad(key_padding_mask, (0, 1))
-        else:
-            assert static_k is None, "bias cannot be added to static key."
-            assert static_v is None, "bias cannot be added to static value."
-    else:
-        assert bias_k is None
-        assert bias_v is None
-
-    q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
-    if k is not None:
-        k = k.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
-    if v is not None:
-        v = v.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
-
-    if static_k is not None:
-        assert static_k.size(0) == bsz * num_heads
-        assert static_k.size(2) == head_dim
-        k = static_k
-
-    if static_v is not None:
-        assert static_v.size(0) == bsz * num_heads
-        assert static_v.size(2) == head_dim
-        v = static_v
-
-    src_len = k.size(1)
-
-    if key_padding_mask is not None:
-        assert key_padding_mask.size(0) == bsz
-        assert key_padding_mask.size(1) == src_len
-
-    if add_zero_attn:
-        src_len += 1
-        k = torch.cat([k, torch.zeros((k.size(0), 1) + k.size()[2:], dtype=k.dtype, device=k.device)], dim=1)
-        v = torch.cat([v, torch.zeros((v.size(0), 1) + v.size()[2:], dtype=v.dtype, device=v.device)], dim=1)
-        if attn_mask is not None:
-            attn_mask = pad(attn_mask, (0, 1))
-        if key_padding_mask is not None:
-            key_padding_mask = pad(key_padding_mask, (0, 1))
-
-    attn_output_weights = torch.bmm(q, k.transpose(1, 2))
-    assert list(attn_output_weights.size()) == [bsz * num_heads, tgt_len, src_len]
-
-    if attn_mask is not None:
-        if attn_mask.dtype == torch.bool:
-            attn_output_weights.masked_fill_(attn_mask, float('-inf'))
-        else:
-            attn_output_weights += attn_mask
-
-
-    if key_padding_mask is not None:
-        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
-        attn_output_weights = attn_output_weights.masked_fill(
-            key_padding_mask.unsqueeze(1).unsqueeze(2),
-            float('-inf'),
-        )
-        attn_output_weights = attn_output_weights.view(bsz * num_heads, tgt_len, src_len)
-
-    attn_output_weights = F.softmax(
-        attn_output_weights, dim=-1)
-    attn_output_weights = F.dropout(attn_output_weights, p=dropout_p, training=training)
-
-    attn_output = torch.bmm(attn_output_weights, v)
-    assert list(attn_output.size()) == [bsz * num_heads, tgt_len, head_dim]
-    attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
-    attn_output = F.linear(attn_output, out_proj_weight, out_proj_bias)
-
-    if need_weights:
-        # average attention weights over heads
-        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
-        return attn_output, attn_output_weights.sum(dim=1) / num_heads
-    else:
-        return attn_output, None
-
-class MultiheadAttention(Module):
-    r"""Allows the model to jointly attend to information
-    from different representation subspaces.
-    See reference: Attention Is All You Need
-    .. math::
-        \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O
-        \text{where} head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)
-    Args:
-        embed_dim: total dimension of the model.
-        num_heads: parallel attention heads.
-        dropout: a Dropout layer on attn_output_weights. Default: 0.0.
-        bias: add bias as module parameter. Default: True.
-        add_bias_kv: add bias to the key and value sequences at dim=0.
-        add_zero_attn: add a new batch of zeros to the key and
-                       value sequences at dim=1.
-        kdim: total number of features in key. Default: None.
-        vdim: total number of features in value. Default: None.
-        Note: if kdim and vdim are None, they will be set to embed_dim such that
-        query, key, and value have the same number of features.
-    Examples::
-        >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
-        >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
-    """
-    # __annotations__ = {
-    #     'bias_k': torch._jit_internal.Optional[torch.Tensor],
-    #     'bias_v': torch._jit_internal.Optional[torch.Tensor],
-    # }
-    __constants__ = ['q_proj_weight', 'k_proj_weight', 'v_proj_weight', 'in_proj_weight']
-
-    def __init__(self, embed_dim, num_heads, dropout=0., bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None):
-        super(MultiheadAttention, self).__init__()
-        self.embed_dim = embed_dim
-        self.kdim = kdim if kdim is not None else embed_dim
-        self.vdim = vdim if vdim is not None else embed_dim
-        self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim
-
-        self.num_heads = num_heads
-        self.dropout = dropout
-        self.head_dim = embed_dim // num_heads
-        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
-
-        if self._qkv_same_embed_dim is False:
-            self.q_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))
-            self.k_proj_weight = Parameter(torch.Tensor(embed_dim, self.kdim))
-            self.v_proj_weight = Parameter(torch.Tensor(embed_dim, self.vdim))
-            self.register_parameter('in_proj_weight', None)
-        else:
-            self.in_proj_weight = Parameter(torch.empty(3 * embed_dim, embed_dim))
-            self.register_parameter('q_proj_weight', None)
-            self.register_parameter('k_proj_weight', None)
-            self.register_parameter('v_proj_weight', None)
-
-        if bias:
-            self.in_proj_bias = Parameter(torch.empty(3 * embed_dim))
-        else:
-            self.register_parameter('in_proj_bias', None)
-        self.out_proj = Linear(embed_dim, embed_dim, bias=bias)
-
-        if add_bias_kv:
-            self.bias_k = Parameter(torch.empty(1, 1, embed_dim))
-            self.bias_v = Parameter(torch.empty(1, 1, embed_dim))
-        else:
-            self.bias_k = self.bias_v = None
-
-        self.add_zero_attn = add_zero_attn
-
-        self._reset_parameters()
-
-    def _reset_parameters(self):
-        if self._qkv_same_embed_dim:
-            xavier_uniform_(self.in_proj_weight)
-        else:
-            xavier_uniform_(self.q_proj_weight)
-            xavier_uniform_(self.k_proj_weight)
-            xavier_uniform_(self.v_proj_weight)
-
-        if self.in_proj_bias is not None:
-            constant_(self.in_proj_bias, 0.)
-            constant_(self.out_proj.bias, 0.)
-        if self.bias_k is not None:
-            xavier_normal_(self.bias_k)
-        if self.bias_v is not None:
-            xavier_normal_(self.bias_v)
-
-    def __setstate__(self, state):
-        # Support loading old MultiheadAttention checkpoints generated by v1.1.0
-        if '_qkv_same_embed_dim' not in state:
-            state['_qkv_same_embed_dim'] = True
-
-        super(MultiheadAttention, self).__setstate__(state)
-
-    def forward(self, query, key, value, key_padding_mask=None,
-                need_weights=True, attn_mask=None):
-        # type: (Tensor, Tensor, Tensor, Optional[Tensor], bool, Optional[Tensor]) -> Tuple[Tensor, Optional[Tensor]]
-        r"""
-    Args:
-        query, key, value: map a query and a set of key-value pairs to an output.
-            See "Attention Is All You Need" for more details.
-        key_padding_mask: if provided, specified padding elements in the key will
-            be ignored by the attention. This is an binary mask. When the value is True,
-            the corresponding value on the attention layer will be filled with -inf.
-        need_weights: output attn_output_weights.
-        attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
-            the batches while a 3D mask allows to specify a different mask for the entries of each batch.
-    Shape:
-        - Inputs:
-        - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
-          the embedding dimension.
-        - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
-          the embedding dimension.
-        - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
-          the embedding dimension.
-        - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
-          If a ByteTensor is provided, the non-zero positions will be ignored while the position
-          with the zero positions will be unchanged. If a BoolTensor is provided, the positions with the
-          value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
-        - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
-          3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
-          S is the source sequence length. attn_mask ensure that position i is allowed to attend the unmasked
-          positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
-          while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
-          is not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
-          is provided, it will be added to the attention weight.
-        - Outputs:
-        - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
-          E is the embedding dimension.
-        - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
-          L is the target sequence length, S is the source sequence length.
-        """
-        if not self._qkv_same_embed_dim:
-            return multi_head_attention_forward(
-                query, key, value, self.embed_dim, self.num_heads,
-                self.in_proj_weight, self.in_proj_bias,
-                self.bias_k, self.bias_v, self.add_zero_attn,
-                self.dropout, self.out_proj.weight, self.out_proj.bias,
-                training=self.training,
-                key_padding_mask=key_padding_mask, need_weights=need_weights,
-                attn_mask=attn_mask, use_separate_proj_weight=True,
-                q_proj_weight=self.q_proj_weight, k_proj_weight=self.k_proj_weight,
-                v_proj_weight=self.v_proj_weight)
-        else:
-            return multi_head_attention_forward(
-                query, key, value, self.embed_dim, self.num_heads,
-                self.in_proj_weight, self.in_proj_bias,
-                self.bias_k, self.bias_v, self.add_zero_attn,
-                self.dropout, self.out_proj.weight, self.out_proj.bias,
-                training=self.training,
-                key_padding_mask=key_padding_mask, need_weights=need_weights,
-                attn_mask=attn_mask)
-
-
-class Transformer(Module):
-    r"""A transformer model. User is able to modify the attributes as needed. The architecture
-    is based on the paper "Attention Is All You Need". Ashish Vaswani, Noam Shazeer,
-    Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and
-    Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information
-    Processing Systems, pages 6000-6010. Users can build the BERT(https://arxiv.org/abs/1810.04805)
-    model with corresponding parameters.
-
-    Args:
-        d_model: the number of expected features in the encoder/decoder inputs (default=512).
-        nhead: the number of heads in the multiheadattention models (default=8).
-        num_encoder_layers: the number of sub-encoder-layers in the encoder (default=6).
-        num_decoder_layers: the number of sub-decoder-layers in the decoder (default=6).
-        dim_feedforward: the dimension of the feedforward network model (default=2048).
-        dropout: the dropout value (default=0.1).
-        activation: the activation function of encoder/decoder intermediate layer, relu or gelu (default=relu).
-        custom_encoder: custom encoder (default=None).
-        custom_decoder: custom decoder (default=None).
-
-    Examples::
-        >>> transformer_model = nn.Transformer(nhead=16, num_encoder_layers=12)
-        >>> src = torch.rand((10, 32, 512))
-        >>> tgt = torch.rand((20, 32, 512))
-        >>> out = transformer_model(src, tgt)
-
-    Note: A full example to apply nn.Transformer module for the word language model is available in
-    https://github.com/pytorch/examples/tree/master/word_language_model
-    """
-
-    def __init__(self, d_model=512, nhead=8, num_encoder_layers=6,
-                 num_decoder_layers=6, dim_feedforward=2048, dropout=0.1,
-                 activation="relu", custom_encoder=None, custom_decoder=None):
-        super(Transformer, self).__init__()
-
-        if custom_encoder is not None:
-            self.encoder = custom_encoder
-        else:
-            encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout, activation)
-            encoder_norm = LayerNorm(d_model)
-            self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
-
-        if custom_decoder is not None:
-            self.decoder = custom_decoder
-        else:
-            decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout, activation)
-            decoder_norm = LayerNorm(d_model)
-            self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm)
-
-        self._reset_parameters()
-
-        self.d_model = d_model
-        self.nhead = nhead
-
-    def forward(self, src, tgt, src_mask=None, tgt_mask=None,
-                memory_mask=None, src_key_padding_mask=None,
-                tgt_key_padding_mask=None, memory_key_padding_mask=None):
-        # type: (Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor]) -> Tensor  # noqa
-        r"""Take in and process masked source/target sequences.
-
-        Args:
-            src: the sequence to the encoder (required).
-            tgt: the sequence to the decoder (required).
-            src_mask: the additive mask for the src sequence (optional).
-            tgt_mask: the additive mask for the tgt sequence (optional).
-            memory_mask: the additive mask for the encoder output (optional).
-            src_key_padding_mask: the ByteTensor mask for src keys per batch (optional).
-            tgt_key_padding_mask: the ByteTensor mask for tgt keys per batch (optional).
-            memory_key_padding_mask: the ByteTensor mask for memory keys per batch (optional).
-
-        Shape:
-            - src: :math:`(S, N, E)`.
-            - tgt: :math:`(T, N, E)`.
-            - src_mask: :math:`(S, S)`.
-            - tgt_mask: :math:`(T, T)`.
-            - memory_mask: :math:`(T, S)`.
-            - src_key_padding_mask: :math:`(N, S)`.
-            - tgt_key_padding_mask: :math:`(N, T)`.
-            - memory_key_padding_mask: :math:`(N, S)`.
-
-            Note: [src/tgt/memory]_mask ensures that position i is allowed to attend the unmasked
-            positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
-            while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
-            are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
-            is provided, it will be added to the attention weight. 
-            [src/tgt/memory]_key_padding_mask provides specified elements in the key to be ignored by
-            the attention. If a ByteTensor is provided, the non-zero positions will be ignored while the zero
-            positions will be unchanged. If a BoolTensor is provided, the positions with the
-            value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
-
-            - output: :math:`(T, N, E)`.
-
-            Note: Due to the multi-head attention architecture in the transformer model,
-            the output sequence length of a transformer is same as the input sequence
-            (i.e. target) length of the decode.
-
-            where S is the source sequence length, T is the target sequence length, N is the
-            batch size, E is the feature number
-
-        Examples:
-            >>> output = transformer_model(src, tgt, src_mask=src_mask, tgt_mask=tgt_mask)
-        """
-
-        if src.size(1) != tgt.size(1):
-            raise RuntimeError("the batch number of src and tgt must be equal")
-
-        if src.size(2) != self.d_model or tgt.size(2) != self.d_model:
-            raise RuntimeError("the feature number of src and tgt must be equal to d_model")
-
-        memory = self.encoder(src, mask=src_mask, src_key_padding_mask=src_key_padding_mask)
-        output = self.decoder(tgt, memory, tgt_mask=tgt_mask, memory_mask=memory_mask,
-                              tgt_key_padding_mask=tgt_key_padding_mask,
-                              memory_key_padding_mask=memory_key_padding_mask)
-        return output
-
-    def generate_square_subsequent_mask(self, sz):
-        r"""Generate a square mask for the sequence. The masked positions are filled with float('-inf').
-            Unmasked positions are filled with float(0.0).
-        """
-        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
-        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
-        return mask
-
-    def _reset_parameters(self):
-        r"""Initiate parameters in the transformer model."""
-
-        for p in self.parameters():
-            if p.dim() > 1:
-                xavier_uniform_(p)
-
-
-class TransformerEncoder(Module):
-    r"""TransformerEncoder is a stack of N encoder layers
-
-    Args:
-        encoder_layer: an instance of the TransformerEncoderLayer() class (required).
-        num_layers: the number of sub-encoder-layers in the encoder (required).
-        norm: the layer normalization component (optional).
-
-    Examples::
-        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
-        >>> transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
-        >>> src = torch.rand(10, 32, 512)
-        >>> out = transformer_encoder(src)
-    """
-    __constants__ = ['norm']
-
-    def __init__(self, encoder_layer, num_layers, norm=None):
-        super(TransformerEncoder, self).__init__()
-        self.layers = _get_clones(encoder_layer, num_layers)
-        self.num_layers = num_layers
-        self.norm = norm
-
-    def forward(self, src, mask=None, src_key_padding_mask=None):
-        # type: (Tensor, Optional[Tensor], Optional[Tensor]) -> Tensor
-        r"""Pass the input through the encoder layers in turn.
-
-        Args:
-            src: the sequence to the encoder (required).
-            mask: the mask for the src sequence (optional).
-            src_key_padding_mask: the mask for the src keys per batch (optional).
-
-        Shape:
-            see the docs in Transformer class.
-        """
-        output = src
-
-        for i, mod in enumerate(self.layers):
-            output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)
-
-        if self.norm is not None:
-            output = self.norm(output)
-
-        return output
-
-
-class TransformerDecoder(Module):
-    r"""TransformerDecoder is a stack of N decoder layers
-
-    Args:
-        decoder_layer: an instance of the TransformerDecoderLayer() class (required).
-        num_layers: the number of sub-decoder-layers in the decoder (required).
-        norm: the layer normalization component (optional).
-
-    Examples::
-        >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
-        >>> transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)
-        >>> memory = torch.rand(10, 32, 512)
-        >>> tgt = torch.rand(20, 32, 512)
-        >>> out = transformer_decoder(tgt, memory)
-    """
-    __constants__ = ['norm']
-
-    def __init__(self, decoder_layer, num_layers, norm=None):
-        super(TransformerDecoder, self).__init__()
-        self.layers = _get_clones(decoder_layer, num_layers)
-        self.num_layers = num_layers
-        self.norm = norm
-
-    def forward(self, tgt, memory, memory2=None, tgt_mask=None,
-                memory_mask=None, memory_mask2=None, tgt_key_padding_mask=None,
-                memory_key_padding_mask=None, memory_key_padding_mask2=None):
-        # type: (Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor]) -> Tensor
-        r"""Pass the inputs (and mask) through the decoder layer in turn.
-
-        Args:
-            tgt: the sequence to the decoder (required).
-            memory: the sequence from the last layer of the encoder (required).
-            tgt_mask: the mask for the tgt sequence (optional).
-            memory_mask: the mask for the memory sequence (optional).
-            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
-            memory_key_padding_mask: the mask for the memory keys per batch (optional).
-
-        Shape:
-            see the docs in Transformer class.
-        """
-        output = tgt
-
-        for mod in self.layers:
-            output = mod(output, memory, memory2=memory2, tgt_mask=tgt_mask,
-                         memory_mask=memory_mask, memory_mask2=memory_mask2,
-                         tgt_key_padding_mask=tgt_key_padding_mask,
-                         memory_key_padding_mask=memory_key_padding_mask,
-                         memory_key_padding_mask2=memory_key_padding_mask2)
-
-        if self.norm is not None:
-            output = self.norm(output)
-
-        return output
-
-class TransformerEncoderLayer(Module):
-    r"""TransformerEncoderLayer is made up of self-attn and feedforward network.
-    This standard encoder layer is based on the paper "Attention Is All You Need".
-    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
-    Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
-    Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
-    in a different way during application.
-
-    Args:
-        d_model: the number of expected features in the input (required).
-        nhead: the number of heads in the multiheadattention models (required).
-        dim_feedforward: the dimension of the feedforward network model (default=2048).
-        dropout: the dropout value (default=0.1).
-        activation: the activation function of intermediate layer, relu or gelu (default=relu).
-
-    Examples::
-        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
-        >>> src = torch.rand(10, 32, 512)
-        >>> out = encoder_layer(src)
-    """
-
-    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, 
-                 activation="relu", debug=False):
-        super(TransformerEncoderLayer, self).__init__()
-        self.debug = debug
-        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
-        # Implementation of Feedforward model
-        self.linear1 = Linear(d_model, dim_feedforward)
-        self.dropout = Dropout(dropout)
-        self.linear2 = Linear(dim_feedforward, d_model)
-
-        self.norm1 = LayerNorm(d_model)
-        self.norm2 = LayerNorm(d_model)
-        self.dropout1 = Dropout(dropout)
-        self.dropout2 = Dropout(dropout)
-
-        self.activation = _get_activation_fn(activation)
-
-    def __setstate__(self, state):
-        if 'activation' not in state:
-            state['activation'] = F.relu
-        super(TransformerEncoderLayer, self).__setstate__(state)
-
-    def forward(self, src, src_mask=None, src_key_padding_mask=None):
-        # type: (Tensor, Optional[Tensor], Optional[Tensor]) -> Tensor
-        r"""Pass the input through the encoder layer.
-
-        Args:
-            src: the sequence to the encoder layer (required).
-            src_mask: the mask for the src sequence (optional).
-            src_key_padding_mask: the mask for the src keys per batch (optional).
-
-        Shape:
-            see the docs in Transformer class.
-        """
-        src2, attn = self.self_attn(src, src, src, attn_mask=src_mask,
-                              key_padding_mask=src_key_padding_mask)
-        if self.debug: self.attn = attn
-        src = src + self.dropout1(src2)
-        src = self.norm1(src)
-        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
-        src = src + self.dropout2(src2)
-        src = self.norm2(src)
-        
-        return src
-
-
-class TransformerDecoderLayer(Module):
-    r"""TransformerDecoderLayer is made up of self-attn, multi-head-attn and feedforward network.
-    This standard decoder layer is based on the paper "Attention Is All You Need".
-    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
-    Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
-    Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
-    in a different way during application.
-
-    Args:
-        d_model: the number of expected features in the input (required).
-        nhead: the number of heads in the multiheadattention models (required).
-        dim_feedforward: the dimension of the feedforward network model (default=2048).
-        dropout: the dropout value (default=0.1).
-        activation: the activation function of intermediate layer, relu or gelu (default=relu).
-
-    Examples::
-        >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
-        >>> memory = torch.rand(10, 32, 512)
-        >>> tgt = torch.rand(20, 32, 512)
-        >>> out = decoder_layer(tgt, memory)
-    """
-
-    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, 
-                 activation="relu", self_attn=True, siamese=False, debug=False):
-        super(TransformerDecoderLayer, self).__init__()
-        self.has_self_attn, self.siamese = self_attn, siamese
-        self.debug = debug
-        if self.has_self_attn:
-            self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
-            self.norm1 = LayerNorm(d_model)
-            self.dropout1 = Dropout(dropout)
-        self.multihead_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
-        # Implementation of Feedforward model
-        self.linear1 = Linear(d_model, dim_feedforward)
-        self.dropout = Dropout(dropout)
-        self.linear2 = Linear(dim_feedforward, d_model)
-
-        self.norm2 = LayerNorm(d_model)
-        self.norm3 = LayerNorm(d_model)
-        self.dropout2 = Dropout(dropout)
-        self.dropout3 = Dropout(dropout)
-        if self.siamese:
-            self.multihead_attn2 = MultiheadAttention(d_model, nhead, dropout=dropout)
-
-        self.activation = _get_activation_fn(activation)
-
-    def __setstate__(self, state):
-        if 'activation' not in state:
-            state['activation'] = F.relu
-        super(TransformerDecoderLayer, self).__setstate__(state)
-
-    def forward(self, tgt, memory, tgt_mask=None, memory_mask=None,
-                tgt_key_padding_mask=None, memory_key_padding_mask=None,
-                memory2=None, memory_mask2=None, memory_key_padding_mask2=None):
-        # type: (Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor]) -> Tensor
-        r"""Pass the inputs (and mask) through the decoder layer.
-
-        Args:
-            tgt: the sequence to the decoder layer (required).
-            memory: the sequence from the last layer of the encoder (required).
-            tgt_mask: the mask for the tgt sequence (optional).
-            memory_mask: the mask for the memory sequence (optional).
-            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
-            memory_key_padding_mask: the mask for the memory keys per batch (optional).
-
-        Shape:
-            see the docs in Transformer class.
-        """
-        if self.has_self_attn:
-            tgt2, attn = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask,
-                                key_padding_mask=tgt_key_padding_mask)
-            tgt = tgt + self.dropout1(tgt2)
-            tgt = self.norm1(tgt)
-            if self.debug: self.attn = attn
-        tgt2, attn2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask,
-                                   key_padding_mask=memory_key_padding_mask)
-        if self.debug: self.attn2 = attn2
-
-        if self.siamese:
-            tgt3, attn3 = self.multihead_attn2(tgt, memory2, memory2, attn_mask=memory_mask2,
-                            key_padding_mask=memory_key_padding_mask2)
-            tgt = tgt + self.dropout2(tgt3)
-            if self.debug: self.attn3 = attn3
-
-        tgt = tgt + self.dropout2(tgt2)
-        tgt = self.norm2(tgt)
-        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
-        tgt = tgt + self.dropout3(tgt2)
-        tgt = self.norm3(tgt)
-        
-        return tgt
-
-
-def _get_clones(module, N):
-    return ModuleList([copy.deepcopy(module) for i in range(N)])
-
-
-def _get_activation_fn(activation):
-    if activation == "relu":
-        return F.relu
-    elif activation == "gelu":
-        return F.gelu
-
-    raise RuntimeError("activation should be relu/gelu, not {}".format(activation))
-
-
-class PositionalEncoding(nn.Module):
-    r"""Inject some information about the relative or absolute position of the tokens
-        in the sequence. The positional encodings have the same dimension as
-        the embeddings, so that the two can be summed. Here, we use sine and cosine
-        functions of different frequencies.
-    .. math::
-        \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))
-        \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))
-        \text{where pos is the word position and i is the embed idx)
-    Args:
-        d_model: the embed dim (required).
-        dropout: the dropout value (default=0.1).
-        max_len: the max. length of the incoming sequence (default=5000).
-    Examples:
-        >>> pos_encoder = PositionalEncoding(d_model)
-    """
-
-    def __init__(self, d_model, dropout=0.1, max_len=5000):
-        super(PositionalEncoding, self).__init__()
-        self.dropout = nn.Dropout(p=dropout)
-
-        pe = torch.zeros(max_len, d_model)
-        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
-        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
-        pe[:, 0::2] = torch.sin(position * div_term)
-        pe[:, 1::2] = torch.cos(position * div_term)
-        pe = pe.unsqueeze(0).transpose(0, 1)
-        self.register_buffer('pe', pe)
-
-    def forward(self, x):
-        r"""Inputs of forward function
-        Args:
-            x: the sequence fed to the positional encoder model (required).
-        Shape:
-            x: [sequence length, batch size, embed dim]
-            output: [sequence length, batch size, embed dim]
-        Examples:
-            >>> output = pos_encoder(x)
-        """
-
-        x = x + self.pe[:x.size(0), :]
-        return self.dropout(x)
-
-
-if __name__ == '__main__':
-    transformer_model = Transformer(nhead=16, num_encoder_layers=12)
-    src = torch.rand((10, 32, 512))
-    tgt = torch.rand((20, 32, 512))
-    out = transformer_model(src, tgt)
-    print(out)

utils.py

@@ -202,7 +202,7 @@ class Config(object):
         assert os.path.exists(config_path), '%s does not exists!' % config_path
         with open(config_path) as file:
             config_dict = yaml.load(file, Loader=yaml.FullLoader)
-        with open('configs/template.yaml') as file:
+        with open('configs/rec/template.yaml') as file:
             default_config_dict = yaml.load(file, Loader=yaml.FullLoader)
         __dict2attr(default_config_dict)
         __dict2attr(config_dict)