OpenMotionLab/MotionGPT

1import numpy as np2import torch3from torch import nn4 5 6class PositionalEncoding(nn.Module):7 8    def __init__(self, d_model, dropout=0.1, max_len=5000, batch_first=False):9        super().__init__()10        self.batch_first = batch_first11 12        self.dropout = nn.Dropout(p=dropout)13 14        pe = torch.zeros(max_len, d_model)15        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)16        div_term = torch.exp(torch.arange(17            0, d_model, 2).float() * (-np.log(10000.0) / d_model))18        pe[:, 0::2] = torch.sin(position * div_term)19        pe[:, 1::2] = torch.cos(position * div_term)20        pe = pe.unsqueeze(0).transpose(0, 1)21 22        self.register_buffer("pe", pe)23 24    def forward(self, x):25        # not used in the final model26        if self.batch_first:27            x = x + self.pe.permute(1, 0, 2)[:, : x.shape[1], :]28        else:29            x = x + self.pe[: x.shape[0], :]30        return self.dropout(x)31