"""
Copyright (c) Meta Platforms, Inc. and affiliates.
All rights reserved.
This source code is licensed under the license found in the
LICENSE file in the root directory of this source tree.
"""
import json
import os
from typing import Any, Dict
import numpy as np
import torch
import torch.optim as optim
from data_loaders.get_data import get_dataset_loader, load_local_data
from diffusion.nn import sum_flat
from model.guide import GuideTransformer
from model.vqvae import setup_tokenizer, TemporalVertexCodec
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
from utils.guide_parser_utils import train_args
from utils.misc import fixseed
class ModelTrainer:
def __init__(
self, args, model: GuideTransformer, tokenizer: TemporalVertexCodec
) -> None:
self.add_frame_cond = args.add_frame_cond
self.data_format = args.data_format
self.tokenizer = tokenizer
self.model = model.cuda()
self.gn = args.gn
self.max_seq_length = args.max_seq_length
self.optimizer = optim.AdamW(
model.parameters(),
lr=args.lr,
betas=(0.9, 0.99),
weight_decay=args.weight_decay,
)
self.scheduler = optim.lr_scheduler.MultiStepLR(
self.optimizer, milestones=args.lr_scheduler, gamma=args.gamma
)
self.l2_loss = lambda a, b: (a - b) ** 2
self.start_step = 0
self.warm_up_iter = args.warm_up_iter
self.lr = args.lr
self.ce_loss = torch.nn.CrossEntropyLoss(
ignore_index=self.tokenizer.n_clusters + 1, label_smoothing=0.1
)
if args.resume_trans is not None:
self._load_from_checkpoint()
def _load_from_checkpoint(self) -> None:
print("loading", args.resume_trans)
ckpt = torch.load(args.resume_trans, map_location="cpu")
self.model.load_state_dict(ckpt["model_state_dict"], strict=True)
self.optimizer.load_state_dict(ckpt["optimizer_state_dict"])
self.start_step = ckpt["iteration"]
def _abbreviate(
self, meshes: torch.Tensor, mask: torch.Tensor, step: int
) -> (torch.Tensor,):
keyframes = meshes[..., ::step]
new_mask = mask[..., ::step]
return keyframes, new_mask
def _prepare_tokens(
self, meshes: torch.Tensor, mask: torch.Tensor
) -> (torch.Tensor,):
if self.add_frame_cond == 1:
keyframes, new_mask = self._abbreviate(meshes, mask, 30)
elif self.add_frame_cond is None:
keyframes, new_mask = self._abbreviate(meshes, mask, 1)
meshes = keyframes.squeeze(2).permute((0, 2, 1))
B, T, _ = meshes.shape
target_tokens = self.tokenizer.predict(meshes)
target_tokens = target_tokens.reshape(B, -1)
input_tokens = torch.cat(
[
torch.zeros(
(B, 1), dtype=target_tokens.dtype, device=target_tokens.device
)
+ self.model.tokens,
target_tokens[:, :-1],
],
axis=-1,
)
return input_tokens, target_tokens, new_mask, meshes.reshape((B, T, -1))
def _run_single_train_step(self, input_tokens, audio, target_tokens):
B, T = input_tokens.shape[0], input_tokens.shape[1]
self.optimizer.zero_grad()
logits = self.model(input_tokens, audio, cond_drop_prob=0.20)
loss = self.ce_loss(
logits.reshape((B * T, -1)), target_tokens.reshape((B * T)).long()
)
loss.backward()
if self.gn:
torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
self.optimizer.step()
self.scheduler.step()
return logits, loss
def _run_single_val_step(
self, motion: torch.Tensor, cond: torch.Tensor
) -> Dict[str, Any]:
self.model.eval()
with torch.no_grad():
motion = torch.as_tensor(motion).cuda()
(
input_tokens,
target_tokens,
new_mask,
downsampled_gt,
) = self._prepare_tokens(motion, cond["mask"])
audio = cond["audio"].cuda()
new_mask = torch.as_tensor(new_mask)
B, T = target_tokens.shape[0], target_tokens.shape[1]
logits = self.model(input_tokens, audio)
tokens = torch.argmax(logits, dim=-1).view(
B, -1, self.tokenizer.residual_depth
)
pred = self.tokenizer.decode(tokens).detach().cpu()
ce_loss = self.ce_loss(
logits.reshape((B * T, -1)), target_tokens.reshape((B * T)).long()
)
l2_loss = self._masked_l2(
downsampled_gt.permute(0, 2, 1).unsqueeze(2).detach().cpu(),
pred.permute(0, 2, 1).unsqueeze(2),
new_mask,
)
acc = self.compute_accuracy(logits, target_tokens, new_mask)
return {
"pred": pred,
"gt": downsampled_gt,
"metrics": {
"ce_loss": ce_loss.item(),
"l2_loss": l2_loss.item(),
"perplexity": np.exp(ce_loss.item()),
"acc": acc.item(),
},
}
def _masked_l2(self, a: torch.Tensor, b: torch.Tensor, mask: torch.Tensor) -> float:
loss = self.l2_loss(a, b)
loss = sum_flat(loss * mask.float())
n_entries = a.shape[1] * a.shape[2]
non_zero_elements = sum_flat(mask) * n_entries
mse_loss_val = loss / non_zero_elements
return mse_loss_val.mean()
def compute_ce_loss(
self, logits: torch.Tensor, target_tokens: torch.Tensor, mask: torch.Tensor
) -> float:
target_tokens[~mask.squeeze().detach().cpu()] = 0
B = logits.shape[0]
logprobs = torch.log_softmax(logits, dim=-1).view(
B, -1, 1, self.tokenizer.n_clusters
)
logprobs = logprobs[:, self.mask_left :, :, :].contiguous()
labels = target_tokens.view(B, -1, 1)
labels = labels[:, self.mask_left :, :].contiguous()
loss = torch.nn.functional.nll_loss(
logprobs.view(-1, self.tokenizer.n_clusters),
labels.view(-1).long(),
reduction="none",
).reshape((B, 1, 1, -1))
mask = mask.float().to(loss.device)
loss = sum_flat(loss * mask)
non_zero_elements = sum_flat(mask)
ce_loss_val = loss / non_zero_elements
return ce_loss_val.mean()
def compute_accuracy(
self, logits: torch.Tensor, target: torch.Tensor, mask: torch.Tensor
) -> float:
mask = mask.squeeze()
probs = torch.softmax(logits, dim=-1)
_, cls_pred_index = torch.max(probs, dim=-1)
acc = (cls_pred_index.flatten(0) == target.flatten(0)).reshape(
cls_pred_index.shape
)
acc = sum_flat(acc).detach().cpu()
non_zero_elements = sum_flat(mask)
acc_val = acc / non_zero_elements * 100
return acc_val.mean()
def update_lr_warm_up(self, nb_iter: int) -> float:
current_lr = self.lr * (nb_iter + 1) / (self.warm_up_iter + 1)
for param_group in self.optimizer.param_groups:
param_group["lr"] = current_lr
return current_lr
def train_step(self, motion: torch.Tensor, cond: torch.Tensor) -> Dict[str, Any]:
self.model.train()
motion = torch.as_tensor(motion).cuda()
input_tokens, target_tokens, new_mask, downsampled_gt = self._prepare_tokens(
motion, cond["mask"]
)
audio = cond["audio"].cuda()
new_mask = torch.as_tensor(new_mask)
logits, loss = self._run_single_train_step(input_tokens, audio, target_tokens)
with torch.no_grad():
tokens = torch.argmax(logits, dim=-1).view(
input_tokens.shape[0], -1, self.tokenizer.residual_depth
)
pred = self.tokenizer.decode(tokens).detach().cpu()
l2_loss = self._masked_l2(
downsampled_gt.permute(0, 2, 1).unsqueeze(2).detach().cpu(),
pred.permute(0, 2, 1).unsqueeze(2),
new_mask,
)
acc = self.compute_accuracy(logits, target_tokens, new_mask)
return {
"pred": pred,
"gt": downsampled_gt,
"loss": loss,
"metrics": {
"ce_loss": loss.item(),
"l2_loss": l2_loss.item(),
"perplexity": np.exp(loss.item()),
"acc": acc.item(),
},
}
def validate(
self,
val_data: DataLoader,
writer: SummaryWriter,
step: int,
save_dir: str,
log_step: int = 100,
max_samples: int = 30,
) -> None:
val_metrics = {}
pred_values = []
gt_values = []
for i, (val_motion, val_cond) in enumerate(val_data):
val_out = self._run_single_val_step(val_motion, val_cond["y"])
if "metrics" in val_out.keys():
for k, v in val_out["metrics"].items():
val_metrics[k] = val_metrics.get(k, 0.0) + v
if "pred" in val_out.keys() and i % log_step == 0:
pred_values.append(
val_data.dataset.inv_transform(val_out["pred"], self.data_format)
)
gt_values.append(
val_data.dataset.inv_transform(val_out["gt"], self.data_format)
)
if i % log_step == 0:
print(
f'val_l2_loss at {step} [{i}]: {val_metrics["l2_loss"] / len(val_data):.4f}'
)
pred_values = torch.concatenate((pred_values), dim=0)
gt_values = torch.concatenate((gt_values), dim=0)
idx = np.random.permutation(len(pred_values))[:max_samples]
pred_values = pred_values[idx]
gt_values = gt_values[idx]
for i, (pred, gt) in enumerate(zip(pred_values, gt_values)):
pred = pred.unsqueeze(0).detach().cpu().numpy()
pose = gt.unsqueeze(0).detach().cpu().numpy()
np.save(os.path.join(save_dir, f"b{i:04d}_pred.npy"), pred)
np.save(os.path.join(save_dir, f"b{i:04d}_gt.npy"), pose)
msg = ""
for k, v in val_metrics.items():
writer.add_scalar(f"val_{k}", v / len(val_data), step)
msg += f"val_{k} at {step}: {v / len(val_data):.4f} | "
print(msg)
def _save_checkpoint(
args, iteration: int, model: GuideTransformer, optimizer: optim.Optimizer
) -> None:
os.makedirs(f"{args.out_dir}/checkpoints/", exist_ok=True)
filename = f"iter-{iteration:07d}.pt"
torch.save(
{
"iteration": iteration,
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
},
f"{args.out_dir}/checkpoints/{filename}",
)
def _load_data_info(args) -> (DataLoader, DataLoader):
data_dict = load_local_data(args.data_root, audio_per_frame=1600)
train_data = get_dataset_loader(
args=args, data_dict=data_dict, split="train", add_padding=False
)
val_data = get_dataset_loader(args=args, data_dict=data_dict, split="val")
return train_data, val_data
def main(args):
fixseed(args.seed)
os.makedirs(args.out_dir, exist_ok=True)
writer = SummaryWriter(f"{args.out_dir}/logs/")
args_path = os.path.join(args.out_dir, "args.json")
with open(args_path, "w") as fw:
json.dump(vars(args), fw, indent=4, sort_keys=True)
tokenizer = setup_tokenizer(args.resume_pth)
model = GuideTransformer(
tokens=tokenizer.n_clusters,
emb_len=798 if args.max_seq_length == 240 else 1998,
num_layers=args.layers,
dim=args.dim,
)
train_data, val_data = _load_data_info(args)
trainer = ModelTrainer(args, model, tokenizer)
step = trainer.start_step
for _ in range(1, args.total_iter + 1):
train_metrics = {}
count = 0
for motion, cond in tqdm(train_data):
if step < args.warm_up_iter:
current_lr = trainer.update_lr_warm_up(step)
# rum single train step
train_out = trainer.train_step(motion, cond["y"])
if "metrics" in train_out.keys():
for k, v in train_out["metrics"].items():
train_metrics[k] = train_metrics.get(k, 0.0) + v
count += 1
# log all of the metrics
if step % args.log_interval == 0:
msg = ""
for k, v in train_metrics.items():
writer.add_scalar(f"train_{k}", v / count, step)
msg += f"train_{k} at {step}: {v / count:.4f} | "
train_metrics = {}
count = 0
writer.add_scalar(f"train_lr", trainer.scheduler.get_lr()[0], step)
if step < args.warm_up_iter:
msg += f"lr: {current_lr} | "
print(msg)
writer.flush()
# run single evaluation step and save
if step % args.eval_interval == 0:
trainer.validate(val_data, writer, step, args.out_dir)
if step % args.save_interval == 0:
_save_checkpoint(args, step, trainer.model, trainer.optimizer)
step += 1
if __name__ == "__main__":
args = train_args()
main(args)