File size: 5,980 Bytes
8b3f571 | 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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | import torch, os
from typing import List, Union
from torch import nn, Tensor
from torch.distributions.distribution import Distribution
from transformers import AutoTokenizer, AutoModel
from typing import Tuple
import numpy as np
from mmengine.registry import MODELS
from mmengine.model import BaseModel
@MODELS.register_module()
class DistilbertEncoderBase(BaseModel):
def __init__(self, modelpath: str,
finetune: bool = False) -> None:
super().__init__()
# Tokenizer
os.environ["TOKENIZERS_PARALLELISM"] = "false"
self.tokenizer = AutoTokenizer.from_pretrained(modelpath)
# Text model
self.text_model = AutoModel.from_pretrained(modelpath)
# Don't train the model
if not finetune:
self.text_model.training = False
for p in self.text_model.parameters():
p.requires_grad = False
# Then configure the model
self.text_encoded_dim = self.text_model.config.hidden_size
def get_last_hidden_state(self, texts: List[str],
return_mask: bool = False
) -> Union[Tensor, Tuple[Tensor, Tensor]]:
encoded_inputs = self.tokenizer(texts, return_tensors="pt", padding=True)
output = self.text_model(**encoded_inputs.to(self.text_model.device))
if not return_mask:
return output.last_hidden_state
return output.last_hidden_state, encoded_inputs.attention_mask.to(dtype=bool)
@MODELS.register_module()
class DistilbertActorAgnosticEncoder(DistilbertEncoderBase):
def __init__(self, modelpath: str,
finetune: bool = False,
vae: bool = True,
latent_dim: int = 256,
ff_size: int = 1024,
num_layers: int = 4, num_heads: int = 4,
dropout: float = 0.1,
activation: str = "gelu", **kwargs) -> None:
super().__init__(modelpath=modelpath, finetune=finetune)
encoded_dim = self.text_encoded_dim
# Projection of the text-outputs into the latent space
self.projection = nn.Sequential(nn.ReLU(),
nn.Linear(encoded_dim, latent_dim))
# TransformerVAE adapted from ACTOR
# Action agnostic: only one set of params
self.vae = vae
if vae:
self.mu_token = nn.Parameter(torch.randn(latent_dim))
self.logvar_token = nn.Parameter(torch.randn(latent_dim))
else:
self.emb_token = nn.Parameter(torch.randn(latent_dim))
self.sequence_pos_encoding = PositionalEncoding(latent_dim, dropout)
seq_trans_encoder_layer = nn.TransformerEncoderLayer(d_model=latent_dim,
nhead=num_heads,
dim_feedforward=ff_size,
dropout=dropout,
activation=activation)
self.seqTransEncoder = nn.TransformerEncoder(seq_trans_encoder_layer,
num_layers=num_layers)
def forward(self, texts: List[str]) -> Union[Tensor, Distribution]:
text_encoded, mask = self.get_last_hidden_state(texts, return_mask=True)
x = self.projection(text_encoded)
bs, nframes, _ = x.shape
# bs, nframes, totjoints, nfeats = x.shape
# Switch sequence and batch_size because the input of
# Pytorch Transformer is [Sequence, Batch size, ...]
x = x.permute(1, 0, 2) # now it is [nframes, bs, latent_dim]
if self.vae:
mu_token = torch.tile(self.mu_token, (bs,)).reshape(bs, -1)
logvar_token = torch.tile(self.logvar_token, (bs,)).reshape(bs, -1)
# adding the distribution tokens for all sequences
xseq = torch.cat((mu_token[None], logvar_token[None], x), 0)
# create a bigger mask, to allow attend to mu and logvar
token_mask = torch.ones((bs, 2), dtype=bool, device=x.device)
aug_mask = torch.cat((token_mask, mask), 1)
else:
emb_token = torch.tile(self.emb_token, (bs,)).reshape(bs, -1)
# adding the embedding token for all sequences
xseq = torch.cat((emb_token[None], x), 0)
# create a bigger mask, to allow attend to emb
token_mask = torch.ones((bs, 1), dtype=bool, device=x.device)
aug_mask = torch.cat((token_mask, mask), 1)
# add positional encoding
xseq = self.sequence_pos_encoding(xseq)
final = self.seqTransEncoder(xseq, src_key_padding_mask=~aug_mask)
if self.vae:
mu, logvar = final[0], final[1]
std = logvar.exp().pow(0.5)
try:
dist = torch.distributions.Normal(mu, std)
except ValueError:
pass
return dist
else:
return final[0]
class PositionalEncoding(nn.Module):
def __init__(self, d_model, dropout=0.1, max_len=5000, batch_first=False):
super().__init__()
self.batch_first = batch_first
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() * (-np.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):
# not used in the final model
if self.batch_first:
x = x + self.pe.permute(1, 0, 2)[:, : x.shape[1], :]
else:
x = x + self.pe[: x.shape[0], :]
return self.dropout(x)
|