Hugging Face's logo

human-motion-capture
/
AnyMo

Model card Files
xet
 Community
AnyMo / src /models /transformers /anymo.py
Xxx999's picture
Xxx999
-
8b3f571
Raw History Blame Contribute Delete
33 kB
import torch, math
import torch.nn.functional as F
from mmengine.registry import MODELS
from mmengine.model import BaseModel
from typing import Dict, List
from torch.distributions.categorical import Categorical
import numpy as np
from random import random
from ...utils.motion import recover_from_local_position_fast
from ..tokenizers.encdoc.encdoc_attn import length_to_mask
def cosine_schedule(t):
return torch.cos(t * math.pi * 0.5)
def get_mask_subset_prob(mask, prob):
subset_mask = torch.bernoulli(mask, p=prob) & mask
return subset_mask
def uniform(shape, device=None):
return torch.zeros(shape, device=device).float().uniform_(0, 1)
def top_k(logits, thres = 0.9, dim = 1):
k = math.ceil((1 - thres) * logits.shape[dim])
val, ind = logits.topk(k, dim = dim)
probs = torch.full_like(logits, float('-inf'))
probs.scatter_(dim, ind, val)
return probs
@MODELS.register_module()
class AnyMo(BaseModel):
def __init__(
self,
vqvae_cfg: Dict,
transformer_cfg: Dict,
max_motion_len: int,
motion_downsample_ratio: int,
audio_downsample_ratio: int,
text_encoder_cfg: Dict=None,
traj_encoder_cfg: Dict=None,
audio_encoder_cfg: Dict=None,
stage=1,
inference_timesteps=6,
cond_types = ['text'],
**kwargs
):
super().__init__(**kwargs)
self.transformer = MODELS.build(transformer_cfg)
# import ipdb; ipdb.set_trace()
if text_encoder_cfg is not None:
self.text_encoder = MODELS.build(text_encoder_cfg)
if traj_encoder_cfg is not None:
self.traj_encoder = MODELS.build(traj_encoder_cfg)
if audio_encoder_cfg is not None:
self.audio_encoder = MODELS.build(audio_encoder_cfg)
# 第一阶段tm训练，第二阶段只训练projector，第三阶段全量训练, 4阶段测评
assert stage in [1, 2, 3, 4]
self.stage = stage
# build vqvae
self.vqvae = MODELS.build(vqvae_cfg)
self.vqvae.init_weights()
for param in self.vqvae.parameters():
param.requires_grad = False
self.vqvae.eval()
self.vocab_size = transformer_cfg.vocab_size # type: ignore
self.n_embd = transformer_cfg.n_embd
self.mask_idx = self.vocab_size - 1
self.pad_idx = self.vocab_size - 2
self.ignore_index = -100
self.max_motion_len = max_motion_len
self.motion_downsample_ratio = motion_downsample_ratio
self.audio_downsample_ratio = audio_downsample_ratio
self.topk_filter_thres = 0.9
self.noise_schedule = cosine_schedule
self.n_quant = len(self.vqvae.quantizer.quantizers)
self.inference_timesteps = inference_timesteps
self.cond_types = cond_types
for cond_type in cond_types:
assert cond_type in ['text', 'audio', 'traj']
if self.stage == 2:
self.init_weights()
# only train audio and traj projector
for param in self.transformer.parameters():
param.requires_grad = False
for param in self.text_encoder.parameters():
param.requires_grad = False
def forward_stage_1(self, motion, motion_length, text_data, mode='loss'):
'''
第一阶段只训t2m
'''
B = motion.shape[0]
# 1. prepare text feature
# text feat: B, text_token_cnt, C
# text mask: B, text_token_cnt
# import ipdb; ipdb.set_trace()
text_feat, text_mask = self.text_encoder(text_data)
audio_feat = torch.zeros(B, 0, self.n_embd).to(text_feat)
audio_mask = torch.zeros(B, 0).to(text_feat)
traj_feat = torch.zeros(B, 0, self.n_embd).to(text_feat)
traj_mask = torch.zeros(B, 0).to(text_feat)
# import ipdb; ipdb.set_trace()
if mode == 'loss':
tokens, label_tokens, non_pad_mask = self.prepare_train_input(motion, motion_length)
loss = self.forward_loss(tokens, label_tokens, text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask, non_pad_mask)
return loss
else:
non_pad_mask = self.prepare_test_input(motion_length)
return self.forward_predict(text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask, non_pad_mask)
def forward_stage_2(self, motion, motion_length, audio_data, traj_data, audio_length, mode='loss', **kwargs):
B = motion.shape[0]
# 2. prepare audio and traj feature
# audio feat: B, audio_token_cnt, C
# traj feat: B, traj_token_cnt, C
prob = random()
# import ipdb; ipdb.set_trace()
if prob < 0.7: # 百分之70的语音数据
audio_feat = self.audio_encoder(audio_data.permute(0, 2, 1)).permute(0, 2, 1)
audio_mask = length_to_mask(audio_length // self.audio_downsample_ratio)
traj_feat = torch.zeros(B, 0, self.n_embd).to(audio_feat)
traj_mask = torch.zeros(B, 0).to(audio_feat)
dtype = audio_feat.dtype
device = audio_feat.device
elif prob >= 0.7 and prob < 0.9: # 百分之20的轨迹数据
traj_feat = self.traj_encoder(traj_data.permute(0, 2, 1)).permute(0, 2, 1)
traj_mask = length_to_mask(motion_length // self.motion_downsample_ratio)
audio_feat = torch.zeros(B, 0, self.n_embd).to(traj_feat)
audio_mask = torch.zeros(B, 0).to(traj_feat)
dtype = traj_feat.dtype
device = traj_feat.device
else:
audio_feat = self.audio_encoder(audio_data.permute(0, 2, 1)).permute(0, 2, 1)
audio_mask = length_to_mask(audio_length // self.audio_downsample_ratio)
traj_feat = self.traj_encoder(traj_data.permute(0, 2, 1)).permute(0, 2, 1)
traj_mask = length_to_mask(motion_length // self.motion_downsample_ratio)
dtype = audio_feat.dtype
device = audio_feat.device
text_feat = torch.zeros(B, 0, self.n_embd).to(dtype).to(device)
text_mask = torch.zeros(B, 0).to(dtype).to(device)
if mode == 'loss':
tokens, label_tokens, non_pad_mask = self.prepare_train_input(motion, motion_length)
loss = self.forward_loss(
tokens, label_tokens, text_feat, text_mask, audio_feat,
audio_mask, traj_feat, traj_mask, non_pad_mask, motion=motion, mean=kwargs['mean'],
std=kwargs['std']
)
return loss
else:
non_pad_mask = self.prepare_test_input(motion_length)
return self.forward_predict(text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask, non_pad_mask, **kwargs)
def forward_stage_3(self, motion, motion_length, text_data, audio_data, traj_data, audio_length, mode='loss', **kwargs):
B = motion.shape[0]
# import ipdb; ipdb.set_trace()
has_traj = random() < 0.1 # 10%的几率加traj
has_text_lst = (audio_length == 0) # 没 audio 的加 text
has_text_lst[has_text_lst == False] = random() < 0.1 # audio 10% 的几率加 text
text_feat, text_mask = self.text_encoder(text_data)
traj_feat = self.traj_encoder(traj_data.permute(0, 2, 1)).permute(0, 2, 1)
traj_mask = length_to_mask(motion_length // self.motion_downsample_ratio)
traj_mask[:] = False if not has_traj else True
text_mask = text_mask * has_text_lst[:, None]
if not torch.all(audio_length == 0):
audio_feat = self.audio_encoder(audio_data.permute(0, 2, 1)).permute(0, 2, 1)
audio_mask = length_to_mask(audio_length // self.audio_downsample_ratio)
else:
# audio下采样8倍
audio_data = torch.zeros((B, 8, 512)).to(motion)
audio_feat = self.audio_encoder(audio_data.permute(0, 2, 1)).permute(0, 2, 1).mean() * 0.0
audio_mask = torch.zeros(B, 0).to(text_feat.device)
text_feat += audio_feat * 0.0
traj_feat += audio_feat * 0.0
audio_feat = torch.zeros(B, 0, self.n_embd).to(text_feat.device)
if mode == 'loss':
tokens, label_tokens, non_pad_mask = self.prepare_train_input(motion, motion_length)
loss = self.forward_loss(
tokens, label_tokens, text_feat, text_mask, audio_feat,
audio_mask, traj_feat, traj_mask, non_pad_mask, motion=motion,
mean=kwargs['mean'], std=kwargs['std']
)
return loss
else:
# 只测text
text_feat, text_mask = self.text_encoder(text_data)
B = text_feat.shape[0]
audio_feat = torch.zeros(B, 0, self.n_embd).to(device)
audio_mask = torch.zeros(B, 0).to(device)
traj_feat = torch.zeros(B, 0, self.n_embd).to(device)
traj_mask = torch.zeros(B, 0).to(device)
non_pad_mask = self.prepare_test_input(motion_length)
return self.forward_predict(text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask, non_pad_mask, **kwargs)
def forward(
self,
motion: torch.Tensor,
motion_length: torch.Tensor,
text_data: List[str]=None,
audio_data: torch.Tensor=None,
traj_data: torch.Tensor=None,
audio_length: torch.Tensor=None,
mode = 'loss',
**kwargs
):
'''
motion: tensor, B, T, C
motion_length: tensor, B
text_data: list[str], len B
has_xxx: list[bool], len B
music_data: tensor, B, T, C
'''
if self.stage == 1:
return self.forward_stage_1(motion, motion_length, text_data, mode)
elif self.stage == 2:
return self.forward_stage_2(motion, motion_length, audio_data, traj_data, audio_length, mode, **kwargs)
elif self.stage == 3:
return self.forward_stage_3(motion, motion_length, text_data, audio_data, traj_data, audio_length, mode, **kwargs)
elif self.stage == 4: # for evaluation
B = motion.shape[0]
device = motion.device
text_feat = torch.zeros(B, 0, self.n_embd).to(device)
audio_feat = torch.zeros(B, 0, self.n_embd).to(device)
traj_feat = torch.zeros(B, 0, self.n_embd).to(device)
text_mask = torch.zeros(B, 0).to(device)
audio_mask = torch.zeros(B, 0).to(device)
traj_mask = torch.zeros(B, 0).to(device)
if 'text' in self.cond_types:
text_feat, text_mask = self.text_encoder(text_data)
dtype = text_feat.dtype
if 'audio' in self.cond_types:
audio_feat = self.audio_encoder(audio_data.permute(0, 2, 1)).permute(0, 2, 1)
audio_mask = length_to_mask(audio_length // self.audio_downsample_ratio)
dtype = audio_feat.dtype
if 'traj' in self.cond_types:
traj_feat = self.traj_encoder(traj_data.permute(0, 2, 1)).permute(0, 2, 1)
traj_mask = length_to_mask(motion_length // self.motion_downsample_ratio)
dtype = traj_feat.dtype
text_feat, text_mask = text_feat.to(dtype), text_mask.to(device)
audio_feat, audio_mask = audio_feat.to(dtype), audio_mask.to(dtype)
traj_feat, traj_mask = traj_feat.to(dtype), traj_mask.to(dtype)
non_pad_mask = self.prepare_test_input(motion_length)
return self.forward_predict(text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask, non_pad_mask, **kwargs)
else:
raise NotImplementedError
def prepare_train_input(self, motion, motion_length):
if self.vqvae.training:
self.vqvae.eval()
B = motion.shape[0]
device = motion.device
m_token_lens = motion_length // self.motion_downsample_ratio
max_token_len = m_token_lens.max()
ids = self.vqvae.encode(motion, motion_length) # B, T, n_quant
# True, True, ..., False
non_pad_mask = torch.arange(max_token_len, device=device).expand(B, max_token_len) < m_token_lens.unsqueeze(1)
all_x_ids = []
all_labels = []
mask = None
for i in range(ids.shape[-1]): # n_quant
# import ipdb; ipdb.set_trace()
m_tokens_i = ids[..., i] # B, T
# 1. prepare label token
labels = torch.where(non_pad_mask, m_tokens_i, self.ignore_index)
# 2. prepare input token
m_tokens_i = torch.where(non_pad_mask, m_tokens_i, self.pad_idx) # [m1, m2, m3, ..., pad1, pad2, ...]
# 2.1 prepare mask
if mask is None:
rand_time = uniform((B,), device=device)
rand_mask_probs = self.noise_schedule(rand_time)
num_token_masked = (max_token_len * rand_mask_probs).round().clamp(min=1)
batch_randperm = torch.rand((B, max_token_len), device=device).argsort(dim=-1)
# Positions to be masked are all true
mask = batch_randperm < num_token_masked.unsqueeze(-1)
# Positions to be MASKED must also be NON-PADDED
mask &= non_pad_mask
# 2.2 Further Apply Bert Masking Scheme
# Step 1: 10% replace with an incorrect token
x_ids = m_tokens_i.clone()
mask_rid = get_mask_subset_prob(mask, 0.1)
rand_id = torch.randint_like(x_ids, high=self.vocab_size - 2)
x_ids = torch.where(mask_rid, rand_id, x_ids)
# Step 2: 90% x 10% replace with correct token, and 90% x 88% replace with mask token
mask_mid = get_mask_subset_prob(mask & ~mask_rid, 0.88)
x_ids = torch.where(mask_mid, self.mask_idx, x_ids)
all_x_ids.append(x_ids)
all_labels.append(labels)
# import ipdb; ipdb.set_trace()
return all_x_ids, all_labels, non_pad_mask
def forward_loss(self,
tokens, label_tokens,
text_feat, text_mask,
audio_feat, audio_mask,
traj_feat, traj_mask,
non_pad_mask, mean=None, std=None, motion=None):
all_logits = self.transformer(tokens, non_pad_mask, text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask)
ce_loss = torch.tensor(0.).to(text_feat)
acc_dict = dict()
for i, (logits, labels) in enumerate(zip(all_logits, label_tokens)):
logits = logits.flatten(0, 1)
labels = labels.flatten(0, 1)
ce_loss = ce_loss + F.cross_entropy(logits, labels.to(torch.int64), ignore_index=self.ignore_index)
# calculate acc
probs = torch.softmax(logits, dim=-1)
cls_pred_index = torch.argmax(probs, dim=-1)
token_mask = (labels != self.ignore_index).int()
right_num = (cls_pred_index == labels).int() * token_mask
acc_dict[f'ACC_{i}'] = (right_num.sum() / token_mask.sum()).float().detach()
# add traj loss
traj_loss = torch.tensor(0.0)
if traj_mask.sum() > 0:
all_pred_motion_tokens = []
for i, (logits, labels) in enumerate(zip(all_logits, label_tokens)):
# import ipdb; ipdb.set_trace()
probs_gumbel_softmax = F.gumbel_softmax(logits, tau=1, eps=1e-10, hard=True, dim = -1)
pred_motion_tokens = torch.topk(probs_gumbel_softmax, k=1, dim=-1)[1].reshape(labels.shape[0], -1) # B, T
all_pred_motion_tokens.append(pred_motion_tokens)
all_pred_motion_tokens = torch.stack(all_pred_motion_tokens).permute(1, 2, 0).to(logits.dtype) # nq, B, T
pred_motion = self.vqvae.decode(all_pred_motion_tokens, motion_length=non_pad_mask.sum(1))
pred_traj = recover_from_local_position_fast(pred_motion.float() * std + mean, njoint=22)[:, :, 0].to(logits.dtype)
gt_traj = recover_from_local_position_fast(motion.float() * std + mean, njoint=22)[:, :, 0].to(logits.dtype)
arange_t = torch.arange(motion.shape[1])[None, :].to(motion.device) # (1, T)
mask = arange_t < (non_pad_mask.sum(1)[:, None] * self.motion_downsample_ratio) # (B, T)
loss_raw = F.l1_loss(pred_traj, gt_traj, reduction='none') # (B, T, C)
traj_loss = (loss_raw * mask[..., None]).sum() / (mask.sum() * loss_raw.shape[-1])
return dict(
ce_loss=ce_loss / self.n_quant,
traj_loss=traj_loss,
**acc_dict
)
@torch.no_grad()
def forward_predict(self, text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask, non_pad_mask, return_ids=False, **kwargs):
device = non_pad_mask.device
# tokens: all pad mask seqence
ids = torch.where(non_pad_mask, self.mask_idx, self.pad_idx)[None].repeat(self.n_quant, 1, 1)
scores = torch.where(non_pad_mask, 0., 1e5)[None].repeat(self.n_quant, 1, 1)
m_lens = non_pad_mask.sum(-1)
for timestep in torch.linspace(0, 1, self.inference_timesteps, device=device):
# import ipdb; ipdb.set_trace()
rand_mask_prob = self.noise_schedule(timestep)
# Maskout, and cope with variable length
num_token_masked = torch.round(rand_mask_prob * m_lens).clamp(min=1) # (b, )
# select num_token_masked tokens with lowest scores to be masked
sorted_indices = scores.argsort(dim=-1) # (b, k), sorted_indices[i, j] = the index of j-th lowest element in scores on dim=1
ranks = sorted_indices.argsort(dim=-1) # (b, k), rank[i, j] = the rank (0: lowest) of scores[i, j] on dim=1
is_mask = (ranks < num_token_masked.unsqueeze(-1).unsqueeze(0))
ids = torch.where(is_mask, self.mask_idx, ids)
# all_logits = self.transformer(ids, non_pad_mask, text_feat, text_mask)
all_logits = self.transformer(ids, non_pad_mask, text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask)
for i, logits in enumerate(all_logits):
# import ipdb; ipdb.set_trace()
filtered_logits = top_k(logits, self.topk_filter_thres, dim=-1)
probs = F.softmax(filtered_logits, dim=-1) # (b, seqlen, ntoken)
pred_ids = Categorical(probs).sample() # (b, seqlen)
ids[i] = torch.where(is_mask[i], pred_ids, ids[i])
# Updating scores
probs_without_temperature = logits.softmax(dim=-1) # (b, seqlen, ntoken)
scores[i] = probs_without_temperature.gather(2, pred_ids.unsqueeze(dim=-1)).squeeze(-1) # (b, seqlen)
# We do not want to re-mask the previously kept tokens, or pad tokens
scores[i] = scores[i].masked_fill(~is_mask[i], 1e5)
ids = torch.where(non_pad_mask[None].repeat(self.n_quant, 1, 1), ids, -1)
# import ipdb; ipdb.set_trace()
pred_motions = []
pred_motion_lengths = []
# import ipdb; ipdb.set_trace()
pred_motions = self.vqvae.decode(ids.permute(1, 2, 0), motion_length=non_pad_mask.sum(1))
if return_ids:
return pred_motions, pred_motion_lengths, ids
else:
return pred_motions
def prepare_test_input(self, motion_length):
B = motion_length.shape[0]
device = motion_length.device
m_lens = motion_length // self.motion_downsample_ratio
max_len = motion_length.max() // self.motion_downsample_ratio
non_pad_mask = torch.arange(max_len, device=device).expand(B, max_len) < m_lens.unsqueeze(1)
return non_pad_mask
@MODELS.register_module()
class AnyMo_audio(BaseModel):
def __init__(
self,
vqvae_cfg: Dict,
transformer_cfg: Dict,
max_motion_len: int,
motion_downsample_ratio: int,
audio_downsample_ratio: int,
text_encoder_cfg: Dict=None,
traj_encoder_cfg: Dict=None,
audio_encoder_cfg: Dict=None,
inference_timesteps=6,
cond_types = ['audio'],
**kwargs
):
super().__init__(**kwargs)
self.transformer = MODELS.build(transformer_cfg)
self.text_encoder = MODELS.build(text_encoder_cfg)
self.traj_encoder = MODELS.build(traj_encoder_cfg)
self.audio_encoder = MODELS.build(audio_encoder_cfg)
# build vqvae
self.vqvae = MODELS.build(vqvae_cfg)
self.vqvae.init_weights()
for param in self.vqvae.parameters():
param.requires_grad = False
self.vqvae.eval()
self.vocab_size = transformer_cfg.vocab_size # type: ignore
self.n_embd = transformer_cfg.n_embd
self.mask_idx = self.vocab_size - 1
self.pad_idx = self.vocab_size - 2
self.ignore_index = -100
self.max_motion_len = max_motion_len
self.motion_downsample_ratio = motion_downsample_ratio
self.audio_downsample_ratio = audio_downsample_ratio
self.topk_filter_thres = 0.9
self.noise_schedule = cosine_schedule
self.n_quant = len(self.vqvae.quantizer.quantizers)
self.inference_timesteps = inference_timesteps
self.cond_types = cond_types
for cond_type in cond_types:
assert cond_type in ['text', 'audio', 'traj']
def forward_train(self, motion, motion_length, text_data, audio_data, traj_data, audio_length, **kwargs):
audio_feat = self.audio_encoder(audio_data.permute(0, 2, 1)).permute(0, 2, 1)
audio_mask = length_to_mask(audio_length // self.audio_downsample_ratio)
B = motion.shape[0]
device = motion.device
text_feat = torch.zeros(B, 0, self.n_embd).to(device).to(audio_feat.dtype)
traj_feat = torch.zeros(B, 0, self.n_embd).to(device).to(audio_feat.dtype)
text_mask = torch.zeros(B, 0).to(device).to(audio_feat.dtype)
traj_mask = torch.zeros(B, 0).to(device).to(audio_feat.dtype)
tokens, label_tokens, non_pad_mask = self.prepare_train_input(motion, motion_length)
loss = self.forward_loss(
tokens, label_tokens, text_feat, text_mask, audio_feat,
audio_mask, traj_feat, traj_mask, non_pad_mask, motion=motion, mean=kwargs['mean'],
std=kwargs['std']
)
return loss
def forward(
self,
motion: torch.Tensor,
motion_length: torch.Tensor,
text_data: List[str]=None,
audio_data: torch.Tensor=None,
traj_data: torch.Tensor=None,
audio_length: torch.Tensor=None,
mode = 'loss',
**kwargs
):
'''
motion: tensor, B, T, C
motion_length: tensor, B
text_data: list[str], len B
has_xxx: list[bool], len B
music_data: tensor, B, T, C
'''
if mode == 'loss':
return self.forward_train(motion, motion_length, text_data, audio_data, traj_data, audio_length, **kwargs)
else: # for evaluation
B = motion.shape[0]
device = motion.device
text_feat = torch.zeros(B, 0, self.n_embd).to(device)
audio_feat = torch.zeros(B, 0, self.n_embd).to(device)
traj_feat = torch.zeros(B, 0, self.n_embd).to(device)
text_mask = torch.zeros(B, 0).to(device)
audio_mask = torch.zeros(B, 0).to(device)
traj_mask = torch.zeros(B, 0).to(device)
# import ipdb; ipdb.set_trace()
if 'text' in self.cond_types:
text_feat, text_mask = self.text_encoder(text_data)
dtype = text_feat.dtype
if 'audio' in self.cond_types:
audio_feat = self.audio_encoder(audio_data.permute(0, 2, 1)).permute(0, 2, 1)
audio_mask = length_to_mask(audio_length // self.audio_downsample_ratio)
dtype = audio_feat.dtype
if 'traj' in self.cond_types:
traj_feat = self.traj_encoder(traj_data.permute(0, 2, 1)).permute(0, 2, 1)
traj_mask = length_to_mask(motion_length // self.motion_downsample_ratio)
dtype = traj_feat.dtype
text_feat, text_mask = text_feat.to(dtype), text_mask.to(device)
audio_feat, audio_mask = audio_feat.to(dtype), audio_mask.to(dtype)
traj_feat, traj_mask = traj_feat.to(dtype), traj_mask.to(dtype)
non_pad_mask = self.prepare_test_input(motion_length)
return self.forward_predict(text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask, non_pad_mask, **kwargs)
def prepare_train_input(self, motion, motion_length):
if self.vqvae.training:
self.vqvae.eval()
B = motion.shape[0]
device = motion.device
m_token_lens = motion_length // self.motion_downsample_ratio
max_token_len = m_token_lens.max()
ids = self.vqvae.encode(motion, motion_length) # B, T, n_quant
# True, True, ..., False
non_pad_mask = torch.arange(max_token_len, device=device).expand(B, max_token_len) < m_token_lens.unsqueeze(1)
all_x_ids = []
all_labels = []
mask = None
for i in range(ids.shape[-1]): # n_quant
# import ipdb; ipdb.set_trace()
m_tokens_i = ids[..., i] # B, T
# 1. prepare label token
labels = torch.where(non_pad_mask, m_tokens_i, self.ignore_index)
# 2. prepare input token
m_tokens_i = torch.where(non_pad_mask, m_tokens_i, self.pad_idx) # [m1, m2, m3, ..., pad1, pad2, ...]
# 2.1 prepare mask
if mask is None:
rand_time = uniform((B,), device=device)
rand_mask_probs = self.noise_schedule(rand_time)
num_token_masked = (max_token_len * rand_mask_probs).round().clamp(min=1)
batch_randperm = torch.rand((B, max_token_len), device=device).argsort(dim=-1)
# Positions to be masked are all true
mask = batch_randperm < num_token_masked.unsqueeze(-1)
# Positions to be MASKED must also be NON-PADDED
mask &= non_pad_mask
# 2.2 Further Apply Bert Masking Scheme
# Step 1: 10% replace with an incorrect token
x_ids = m_tokens_i.clone()
mask_rid = get_mask_subset_prob(mask, 0.1)
rand_id = torch.randint_like(x_ids, high=self.vocab_size - 2)
x_ids = torch.where(mask_rid, rand_id, x_ids)
# Step 2: 90% x 10% replace with correct token, and 90% x 88% replace with mask token
mask_mid = get_mask_subset_prob(mask & ~mask_rid, 0.88)
x_ids = torch.where(mask_mid, self.mask_idx, x_ids)
all_x_ids.append(x_ids)
all_labels.append(labels)
# import ipdb; ipdb.set_trace()
return all_x_ids, all_labels, non_pad_mask
def forward_loss(self,
tokens, label_tokens,
text_feat, text_mask,
audio_feat, audio_mask,
traj_feat, traj_mask,
non_pad_mask, mean=None, std=None, motion=None):
all_logits = self.transformer(tokens, non_pad_mask, text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask)
ce_loss = torch.tensor(0.).to(text_feat)
acc_dict = dict()
for i, (logits, labels) in enumerate(zip(all_logits, label_tokens)):
logits = logits.flatten(0, 1)
labels = labels.flatten(0, 1)
ce_loss = ce_loss + F.cross_entropy(logits, labels.to(torch.int64), ignore_index=self.ignore_index)
# calculate acc
probs = torch.softmax(logits, dim=-1)
cls_pred_index = torch.argmax(probs, dim=-1)
token_mask = (labels != self.ignore_index).int()
right_num = (cls_pred_index == labels).int() * token_mask
acc_dict[f'ACC_{i}'] = (right_num.sum() / token_mask.sum()).float().detach()
# add traj loss
traj_loss = torch.tensor(0.0)
if traj_mask.sum() > 0:
all_pred_motion_tokens = []
for i, (logits, labels) in enumerate(zip(all_logits, label_tokens)):
# import ipdb; ipdb.set_trace()
probs_gumbel_softmax = F.gumbel_softmax(logits, tau=1, eps=1e-10, hard=True, dim = -1)
pred_motion_tokens = torch.topk(probs_gumbel_softmax, k=1, dim=-1)[1].reshape(labels.shape[0], -1) # B, T
all_pred_motion_tokens.append(pred_motion_tokens)
all_pred_motion_tokens = torch.stack(all_pred_motion_tokens) # nq, B, T
pred_motion = self.vqvae.decode(all_pred_motion_tokens.permute(1, 2, 0), motion_length=non_pad_mask.sum(1))
pred_traj = recover_from_local_position_fast(pred_motion * std + mean, njoint=22)[:, :, 0]
gt_traj = recover_from_local_position_fast(motion * std + mean, njoint=22)[:, :, 0]
arange_t = torch.arange(motion.shape[1])[None, :].to(motion.device) # (1, T)
mask = arange_t < (non_pad_mask.sum(1)[:, None] * self.motion_downsample_ratio) # (B, T)
loss_raw = F.l1_loss(pred_traj, gt_traj, reduction='none') # (B, T, C)
traj_loss = (loss_raw * mask[..., None]).sum() / (mask.sum() * loss_raw.shape[-1])
return dict(
ce_loss=ce_loss / self.n_quant,
traj_loss=traj_loss,
**acc_dict
)
@torch.no_grad()
def forward_predict(self, text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask, non_pad_mask, return_ids=False, **kwargs):
device = non_pad_mask.device
# tokens: all pad mask seqence
ids = torch.where(non_pad_mask, self.mask_idx, self.pad_idx)[None].repeat(self.n_quant, 1, 1)
scores = torch.where(non_pad_mask, 0., 1e5)[None].repeat(self.n_quant, 1, 1)
m_lens = non_pad_mask.sum(-1)
for timestep in torch.linspace(0, 1, self.inference_timesteps, device=device):
# import ipdb; ipdb.set_trace()
rand_mask_prob = self.noise_schedule(timestep)
# Maskout, and cope with variable length
num_token_masked = torch.round(rand_mask_prob * m_lens).clamp(min=1) # (b, )
# select num_token_masked tokens with lowest scores to be masked
sorted_indices = scores.argsort(dim=-1) # (b, k), sorted_indices[i, j] = the index of j-th lowest element in scores on dim=1
ranks = sorted_indices.argsort(dim=-1) # (b, k), rank[i, j] = the rank (0: lowest) of scores[i, j] on dim=1
is_mask = (ranks < num_token_masked.unsqueeze(-1).unsqueeze(0))
ids = torch.where(is_mask, self.mask_idx, ids)
# all_logits = self.transformer(ids, non_pad_mask, text_feat, text_mask)
all_logits = self.transformer(ids, non_pad_mask, text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask)
for i, logits in enumerate(all_logits):
# import ipdb; ipdb.set_trace()
filtered_logits = top_k(logits, self.topk_filter_thres, dim=-1)
probs = F.softmax(filtered_logits, dim=-1) # (b, seqlen, ntoken)
pred_ids = Categorical(probs).sample() # (b, seqlen)
ids[i] = torch.where(is_mask[i], pred_ids, ids[i])
# Updating scores
probs_without_temperature = logits.softmax(dim=-1) # (b, seqlen, ntoken)
scores[i] = probs_without_temperature.gather(2, pred_ids.unsqueeze(dim=-1)).squeeze(-1) # (b, seqlen)
# We do not want to re-mask the previously kept tokens, or pad tokens
scores[i] = scores[i].masked_fill(~is_mask[i], 1e5)
ids = torch.where(non_pad_mask[None].repeat(self.n_quant, 1, 1), ids, -1)
# import ipdb; ipdb.set_trace()
pred_motions = []
pred_motion_lengths = []
# import ipdb; ipdb.set_trace()
pred_motions = self.vqvae.decode(ids.permute(1, 2, 0), motion_length=non_pad_mask.sum(1))
if return_ids:
return pred_motions, pred_motion_lengths, ids
else:
return pred_motions
def prepare_test_input(self, motion_length):
B = motion_length.shape[0]
device = motion_length.device
m_lens = motion_length // self.motion_downsample_ratio
max_len = motion_length.max() // self.motion_downsample_ratio
non_pad_mask = torch.arange(max_len, device=device).expand(B, max_len) < m_lens.unsqueeze(1)
return non_pad_mask