vumichien/Generate_human_motion

1"""2Various positional encodings for the transformer.3"""4import math5import torch6from torch import nn7 8def PE1d_sincos(seq_length, dim):9    """10    :param d_model: dimension of the model11    :param length: length of positions12    :return: length*d_model position matrix13    """14    if dim % 2 != 0:15        raise ValueError("Cannot use sin/cos positional encoding with "16                         "odd dim (got dim={:d})".format(dim))17    pe = torch.zeros(seq_length, dim)18    position = torch.arange(0, seq_length).unsqueeze(1)19    div_term = torch.exp((torch.arange(0, dim, 2, dtype=torch.float) *20                         -(math.log(10000.0) / dim)))21    pe[:, 0::2] = torch.sin(position.float() * div_term)22    pe[:, 1::2] = torch.cos(position.float() * div_term)23 24    return pe.unsqueeze(1)25 26 27class PositionEmbedding(nn.Module):28    """29    Absolute pos embedding (standard), learned.30    """31    def __init__(self, seq_length, dim, dropout, grad=False):32        super().__init__()33        self.embed = nn.Parameter(data=PE1d_sincos(seq_length, dim), requires_grad=grad)34        self.dropout = nn.Dropout(p=dropout)35        36    def forward(self, x):37        # x.shape: bs, seq_len, feat_dim38        l = x.shape[1]39        x = x.permute(1, 0, 2) + self.embed[:l].expand(x.permute(1, 0, 2).shape)40        x = self.dropout(x.permute(1, 0, 2))41        return x42 43