Spaces:
Build error
Build error
File size: 3,186 Bytes
cb433d6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
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
|