operator optimization

[boji123]

https://github.com/ubisoft/ubisoft-laforge-ZeroEGGS/blob/4992d9ab5e28bc4bb13219a598906562f5786448/ZEGGS/modules.py#L461C5-L481

consider relace with

def forward(self, x):
  """ Forward function of Positional Encoding:
      x = (B, N) -- Long or Int tensor
  """
  # initialize tensor
  nb_frames_max = torch.max(torch.cumsum(x, dim=1))
  self.pos_enc = self.pos_enc.to(x.device)
  
  # [nb_frames_max]
  seq_range = torch.arange(0, nb_frames_max, dtype=torch.int64, device=x.device)
  # [batch_size, nb_frames_max]
  seq_range_expand = seq_range.unsqueeze(0).expand(x.size(0), nb_frames_max)
  seq_mask = seq_range_expand < x
  # [batch_size, nb_frames_max, embed_dim]
  pos_emb = self.pos_enc.index_select(0, seq_range_expand.flatten()).reshape(x.size(0), nb_frames_max, -1)
  pos_emb = pos_emb * seq_mask.unsqueeze(2)
  
  return pos_emb

[MengHao666]

woud it affect the precision ? could your method only speed up the code, but would not improve precision in any means?

[boji123]

only speed up, friendly to jit export

[MengHao666]

only speed up, friendly to jit export

I have a question about the data preprocess of the code. I am now try to use such method to train on BEAT dataset. Could I add your wechat to ask some questions? Best wishes. My wechat id is 19800361080 .