class TransformerModel(nn.Module): def __init__(self, dim_in, units, nhead, dim_out, nlayers, dropout=0.5): super(TransformerModel, self).__init__() from torch.nn import TransformerEncoder, TransformerEncoderLayer self.emb = nn.Linear(dim_in, units) encoder_layers = TransformerEncoderLayer(units, nhead) self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers) self.decoder = nn.Linear(units, dim_out) self.optimizer = optim.Adam(chain( self.transformer_encoder.parameters(), self.emb.parameters(), self.decoder.parameters()), lr=0.01, betas=(0.5, 0.999)) self.cuda = False if not torch.cuda.is_available() else True self.activation = F.relu self.dropout = nn.Dropout(dropout) self.Tensor = torch.cuda.FloatTensor if self.cuda else torch.Tensor self.softmax = nn.LogSoftmax(dim=1) # self.loss_fn = EMDLoss() self.loss_fn = nn.CrossEntropyLoss() def forward(self, src): src = self.emb(src) encoder_output = self.transformer_encoder(self.activation(src)) output = self.decoder(encoder_output.mean(dim=1)) # output = self.decoder(encoder_output[:,0]) return output def train(self, batch): label = batch["label"] if self.cuda: target = move_to_gpu(label) else: target = label src = Variable(batch['x'].type(self.Tensor)) logits = self.forward(src) loss = self.loss_fn(logits, target) self.optimizer.zero_grad() # logits.retain_grad() loss.backward() # print("grad", logits.grad) self.optimizer.step() metrics = { "loss": loss.item(), } return metrics def predict(self, batch): src = Variable(batch['x'].type(self.Tensor)) logits = self.forward(src) scores = self.softmax(logits) scores = torch.exp(scores) return scores
class Transformer(nn.Module): def __init__(self, vocab_size: int, d_model: int, nhead: int, d_feedforward: int, nlayers: int, dropout: float = 0.5, max_len: int = 5, norm_first=False, use_norm=False, **kwargs): super().__init__() self.vocab_size = vocab_size self.model_type = 'TransformerDecoder' self.target_pad_idx = 3 self.pos_encoder = PositionalEncoding(d_model, dropout, max_len=max_len) transformer_layers = TransformerEncoderLayer( d_model, nhead, dim_feedforward=d_feedforward, dropout=dropout, norm_first=norm_first, **kwargs) if use_norm: transformer_norm = LayerNorm(d_model, eps=1e-5, **kwargs) else: transformer_norm = None self.transformer_block = TransformerEncoder(transformer_layers, nlayers, transformer_norm) self.encoder = nn.Embedding(vocab_size, d_model) self.d_model = d_model self.decoder = nn.Linear(d_model, vocab_size) self.max_len = max_len self.criterion = nn.CrossEntropyLoss(ignore_index=self.target_pad_idx) self._reset_parameters() self.init_weights() self.input_attn_mask = self.generate_square_subsequent_mask( self.max_len) @staticmethod def generate_square_subsequent_mask(sz: int) -> Tensor: r"""Generate a square mask for the sequence. The masked positions are filled with float('-inf'). Unmasked positions are filled with float(0.0). """ return torch.triu(torch.full((sz, sz), float('-inf')), diagonal=1) def _reset_parameters(self): r"""Initiate parameters in the transformer model.""" for p in self.transformer_block.parameters(): if p.dim() > 1: xavier_uniform_(p) def init_weights(self) -> None: initrange = 0.1 self.encoder.weight.data.uniform_(-initrange, initrange) self.decoder.bias.data.zero_() self.decoder.weight.data.uniform_(-initrange, initrange) def forward(self, x: Tensor, x_lens=None) -> Tensor: """ Args: src: Tensor, shape [seq_len, batch_size] src_mask: Tensor, shape [seq_len, seq_len] Returns: output Tensor of shape [seq_len, batch_size, vocab_size] """ x = x.transpose(0, 1) src = self.encoder(x) * math.sqrt(self.d_model) # src = self.pos_encoder(src) self.input_attn_mask = self.generate_square_subsequent_mask( src.shape[0]) output = self.transformer_block(src, self.input_attn_mask) output = self.decoder(output) # pred = output[-1] # torch.softmax(output[-1], 1) return output.transpose(0, 1).contiguous() def compute_loss(self, preds, targets): targets = targets.view(-1, self.vocab_size) preds = preds.view(-1, self.vocab_size) loss = self.criterion(preds, targets.argmax(dim=1)) return loss
class MusicBert(torch.nn.Module): def __init__(self, dim_sound, num_tokens=30522, dim_model=768, nhead=12, name="music_bert",\ num_encoder_layers=12, d_feed=3072, dropout=0.1, n_convs=4): super(MusicBert, self).__init__() self.d_model = dim_model self.num_tokens = num_tokens self.PATH = "models/" + name + ".pth" # Token Embedder self.embedder = nn.Embedding(num_tokens, dim_model) self.sound_compressor = SoundCNNEncoder([dim_sound] + [dim_model] * (n_convs - 1)) self.sound_decompressor = SoundCNNDecoder([dim_model] * n_convs) self.position_embeddings = PositionalEncodings(dim_model, dropout) encoder_layer = TransformerEncoderLayer(dim_model, nhead, \ dim_feedforward=d_feed, \ dropout=dropout, activation='gelu') device = self.get_device() # Registers IDs in the buffers memory self.register_buffer("CLS_ID", torch.tensor([[101]]).to(device).long()) self.register_buffer("SEP_ID", torch.tensor([[102]]).to(device).long()) self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers) def set_encoder_grad(self, grad): for param in self.encoder.parameters(): param.requires_grad = grad def set_convs_grad(self, grad): for param in self.sound_compressor.parameters(): param.requires_grad = grad for param in self.sound_decompressor.parameters(): param.requires_grad = grad def init_weights(self): initrange = 0.1 self.encoder.weight.data.uniform_(-initrange, initrange) def get_device(self): return next(self.parameters()).device def load_pretrained(self): #load pretrained model from Hugginface device = self.get_device() bertModel = BertModel.from_pretrained("bert-base-uncased").to(device) bert_state_dict = translate_from_hugginface_to_torch_BERT( bertModel.state_dict()) self.load_state_dict(bert_state_dict, strict=False) self.__test_integrity() def __test_integrity(self): self.eval() tokenizer = BertTokenizer.from_pretrained('bert-base-uncased') BERTmodel = BertModel.from_pretrained('bert-base-uncased').to( self.get_device()) input_ids = torch.tensor( tokenizer.encode("this is a test")).unsqueeze(0).to( self.get_device()) BERT_outputs = BERTmodel(input_ids)[0] my_outputs = self(src_tokens=input_ids.permute(1, 0), add_special_tokens=False).permute(1, 0, 2) outputs_mean = torch.mean(my_outputs - BERT_outputs) outputs_norm = torch.norm(my_outputs - BERT_outputs) print("== INTEGRITY TEST ==") print(f"Mean error: {outputs_mean}") print(f"Norm error: {outputs_norm}") self.train() def save_model(self, path=None): path = path if path is not None else self.PATH torch.save(self.state_dict(), path) def load_model(self, path=None): path = path if path is not None else self.PATH self.load_state_dict(torch.load(path)) def forward(self, src_sound=None, src_tokens=None, add_special_tokens=True, skip_encoder=False): assert src_tokens is not None or src_sound is not None, \ "Feed at least one sound or token sequence" # SHAPES: sound Seq_len, Batch, Feat # tokens Seq_len, Batch sound_length = 0 tokens_length = 0 ## "Preprocess" the two sequences to the same Vector space if src_sound is not None: batch_size = src_sound.shape[1] sound_length = src_sound.shape[0] src_sound = self.sound_compressor(src_sound) src_sound = self.position_embeddings(src_sound, token_type="sound") else: # if undefined creates an empty placeholder for the following operations src_sound = torch.Tensor([]).to(self.get_device()) if src_tokens is not None: batch_size = src_tokens.shape[1] tokens_length = src_tokens.shape[0] else: # if undefined creates an empty placeholder for the following operations src_tokens = torch.Tensor([]).to(self.get_device()) if add_special_tokens: CLS = self.CLS_ID.repeat(1, batch_size) SEP = self.SEP_ID.repeat(1, batch_size) src_tokens = torch.cat((CLS, src_tokens.long(), SEP), 0) tokens_length += 2 # we added 2 tokens in the token_sequence if tokens_length != 0: src_tokens = self.embedder(src_tokens.long()) src_tokens = self.position_embeddings(src_tokens, token_type="text") sequence = torch.cat((src_tokens, src_sound), 0) if skip_encoder: output = sequence else: output = self.encoder(sequence) if sound_length > 0: # if sound was provided, decompresses the sound_sequence sound_sequence = output[tokens_length:] token_sequence = output[:tokens_length] sound_sequence = self.sound_decompressor(sound_sequence, target_len=sound_length) output = torch.cat((token_sequence, sound_sequence), 0) assert output.shape[0] == tokens_length + sound_length,\ "Expecting that input and output sequences to have the same length, but got {:d} and {:d}"\ .format(output.shape[0], tokens_length + sound_length) return output